diff --git a/Cargo.lock b/Cargo.lock
index 905235f4..838be979 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2498,12 +2498,14 @@ dependencies = [
  "rustrain-checkpoint",
  "rustrain-core",
  "rustrain-ipc",
+ "rustrain-parallel",
  "rustrain-qwen3-6",
  "rustrain-train",
  "safetensors 0.8.0",
  "serde",
  "serde_json",
  "tch",
+ "tempfile",
  "tokenizers",
  "tokio",
  "tokio-stream",
diff --git a/crates/rustrain-core/src/runtime.rs b/crates/rustrain-core/src/runtime.rs
index eb19038f..30a0e093 100644
--- a/crates/rustrain-core/src/runtime.rs
+++ b/crates/rustrain-core/src/runtime.rs
@@ -427,6 +427,23 @@ pub fn load_config(path: &Path) -> Result<Config> {
     toml::from_str(&contents).with_context(|| format!("failed to parse config {}", path.display()))
 }
 
+fn qwen_source_parallel_factor(
+    is_qwen_hybrid: bool,
+    is_expert_parallel_architecture: bool,
+    ep_source_sharded: bool,
+    parallel: &ParallelConfig,
+) -> usize {
+    if !is_qwen_hybrid {
+        return 1;
+    }
+    let ep_source_factor = if is_expert_parallel_architecture && ep_source_sharded {
+        parallel.expert_model_parallel_size
+    } else {
+        1
+    };
+    parallel.data_parallel_size * ep_source_factor
+}
+
 pub fn validate_config(config: &Config) -> Result<()> {
     if matches!(config.train.backend, BackendKind::NdArray)
         && !matches!(config.train.device, Device::Cpu)
@@ -512,6 +529,11 @@ pub fn validate_config(config: &Config) -> Result<()> {
         && config.model.architecture == "qwen_trainable_session";
     let is_qwen_lora_sft = matches!(config.train.backend, BackendKind::Tch)
         && config.model.architecture == "qwen_lora_sft";
+    let is_qwen3_hybrid_lora_sft = matches!(config.train.backend, BackendKind::Tch)
+        && matches!(
+            config.model.architecture.as_str(),
+            "qwen3_5_lora_sft" | "qwen3_5_lora_sft_ep" | "qwen3_6_lora_sft" | "qwen3_6_lora_sft_ep"
+        );
     let is_tch_moe_ep_session = matches!(config.train.backend, BackendKind::Tch)
         && config.model.architecture == "tch_moe_ep_session";
     let is_v4_tp_rank = matches!(config.train.backend, BackendKind::Tch)
@@ -526,14 +548,23 @@ pub fn validate_config(config: &Config) -> Result<()> {
         && config.model.architecture == "deepseek_v4_lora_sft_ep";
     let is_glm5_lora_sft_ep = matches!(config.train.backend, BackendKind::Tch)
         && config.model.architecture == "glm5_lora_sft_ep";
-    let is_qwen3_6_lora_sft_ep = matches!(config.train.backend, BackendKind::Tch)
-        && config.model.architecture == "qwen3_6_lora_sft_ep";
+    let is_qwen3_hybrid_lora_sft_ep = matches!(config.train.backend, BackendKind::Tch)
+        && matches!(
+            config.model.architecture.as_str(),
+            "qwen3_5_lora_sft_ep" | "qwen3_6_lora_sft_ep"
+        );
     let is_v4_tp_ep_train = matches!(config.train.backend, BackendKind::Tch)
         && config.model.architecture == "deepseek_v4_tp_ep_train";
     let is_v3_tp_rank = matches!(config.train.backend, BackendKind::Tch)
         && config.model.architecture == "deepseek_tp_rank";
     let is_v3_ep_rank = matches!(config.train.backend, BackendKind::Tch)
         && config.model.architecture == "deepseek_ep_rank";
+    if is_qwen3_hybrid_lora_sft_ep && parallel.context_parallel_size != 1 {
+        return Err(anyhow!(
+            "{} does not yet support context parallelism; context_parallel_size must be 1",
+            config.model.architecture
+        ));
+    }
     for (name, value) in parallel_sizes {
         if value == 0 {
             return Err(anyhow!("{name} must be greater than zero"));
@@ -542,10 +573,23 @@ pub fn validate_config(config: &Config) -> Result<()> {
             && !((is_tch_tiny_lm || is_qwen_trainable_session)
                 && name == "data_parallel_size"
                 && value == 2)
+            && !(is_qwen3_hybrid_lora_sft
+                && name == "data_parallel_size"
+                && value >= 2
+                && parallel.context_parallel_size == 1)
+            && !(is_qwen3_hybrid_lora_sft
+                && name == "tensor_model_parallel_size"
+                && parallel.context_parallel_size == 1)
+            && !(is_qwen3_hybrid_lora_sft
+                && name == "pipeline_model_parallel_size"
+                && parallel.context_parallel_size == 1)
             && !(is_qwen_trainable_session
                 && name == "tensor_model_parallel_size"
                 && value == 2
                 && parallel.data_parallel_size == 1)
+            && !(is_qwen3_hybrid_lora_sft_ep
+                && name == "expert_model_parallel_size"
+                && parallel.context_parallel_size == 1)
             && !(is_tch_moe_ep_session
                 && name == "expert_model_parallel_size"
                 && value == 2
@@ -559,7 +603,7 @@ pub fn validate_config(config: &Config) -> Result<()> {
                 || is_v4_ep_train
                 || is_v4_lora_sft_ep
                 || is_glm5_lora_sft_ep
-                || is_qwen3_6_lora_sft_ep)
+                || is_qwen3_hybrid_lora_sft_ep)
                 && name == "expert_model_parallel_size"
                 && parallel.data_parallel_size == 1)
             && !(is_glm5_lora_sft_ep
@@ -573,7 +617,7 @@ pub fn validate_config(config: &Config) -> Result<()> {
         }
     }
 
-    if is_qwen_trainable_session || is_qwen_lora_sft {
+    if is_qwen_trainable_session || is_qwen_lora_sft || is_qwen3_hybrid_lora_sft {
         if !matches!(config.train.device, Device::Cuda) {
             return Err(anyhow!(
                 "{} requires device = \"cuda\"",
@@ -608,7 +652,7 @@ pub fn validate_config(config: &Config) -> Result<()> {
         }
     }
 
-    if is_qwen_lora_sft {
+    if is_qwen_lora_sft || is_qwen3_hybrid_lora_sft {
         let lora = config
             .lora
             .as_ref()
@@ -634,20 +678,44 @@ pub fn validate_config(config: &Config) -> Result<()> {
         if lora.target_modules.is_empty() {
             return Err(anyhow!("lora.target_modules must not be empty"));
         }
-        let supported_lora_modules = [
-            "q_proj",
-            "k_proj",
-            "v_proj",
-            "o_proj",
-            "gate_proj",
-            "up_proj",
-            "down_proj",
-        ];
+        let supported_lora_modules: &[&str] = if is_qwen3_hybrid_lora_sft {
+            &[
+                "q_proj",
+                "k_proj",
+                "v_proj",
+                "o_proj",
+                "in_proj_qkv",
+                "in_proj_z",
+                "in_proj_a",
+                "in_proj_b",
+                "out_proj",
+                "gate_proj",
+                "up_proj",
+                "down_proj",
+                "shared_gate_proj",
+                "shared_up_proj",
+                "shared_down_proj",
+                "experts_gate_up_proj",
+                "experts_down_proj",
+            ]
+        } else {
+            &[
+                "q_proj",
+                "k_proj",
+                "v_proj",
+                "o_proj",
+                "gate_proj",
+                "up_proj",
+                "down_proj",
+            ]
+        };
         for module in &lora.target_modules {
             if !supported_lora_modules.contains(&module.as_str()) {
                 return Err(anyhow!(
-                    "qwen_lora_sft unsupported lora.target_modules entry {}; supported: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj",
-                    module
+                    "{} unsupported lora.target_modules entry {}; supported: {}",
+                    config.model.architecture,
+                    module,
+                    supported_lora_modules.join(", ")
                 ));
             }
         }
@@ -658,17 +726,38 @@ pub fn validate_config(config: &Config) -> Result<()> {
             .any(|module| !unique_modules.insert(module))
         {
             return Err(anyhow!(
-                "qwen_lora_sft lora.target_modules must not contain duplicates"
+                "{} lora.target_modules must not contain duplicates",
+                config.model.architecture
             ));
         }
         if config.train.micro_batch_size == 0 {
-            return Err(anyhow!("qwen_lora_sft requires micro_batch_size > 0"));
+            return Err(anyhow!(
+                "{} requires micro_batch_size > 0",
+                config.model.architecture
+            ));
         }
-        let expected_global_batch_size =
-            config.train.micro_batch_size * config.train.gradient_accumulation_steps;
+        let ep_source_sharded = std::env::var("QWEN36_EP_A2A_SHARDED")
+            .map(|value| !value.is_empty() && value != "0")
+            .unwrap_or(false);
+        let source_parallel_factor = qwen_source_parallel_factor(
+            is_qwen3_hybrid_lora_sft,
+            is_qwen3_hybrid_lora_sft_ep,
+            ep_source_sharded,
+            &config.parallel,
+        );
+        let expected_global_batch_size = config.train.micro_batch_size
+            * config.train.gradient_accumulation_steps
+            * source_parallel_factor;
         if config.train.global_batch_size != expected_global_batch_size {
+            if source_parallel_factor > 1 {
+                return Err(anyhow!(
+                    "{} requires global_batch_size = micro_batch_size * gradient_accumulation_steps * effective source-parallel size ({source_parallel_factor}); this includes data_parallel_size and, for sharded EP sources, expert_model_parallel_size",
+                    config.model.architecture,
+                ));
+            }
             return Err(anyhow!(
-                "qwen_lora_sft requires global_batch_size = micro_batch_size * gradient_accumulation_steps"
+                "{} requires global_batch_size = micro_batch_size * gradient_accumulation_steps",
+                config.model.architecture
             ));
         }
     } else if config.train.micro_batch_size != 1 || config.train.global_batch_size != 1 {
@@ -1207,8 +1296,8 @@ pub fn validate_config(config: &Config) -> Result<()> {
 }
 
 pub fn prepare_run_directory(run: &RunConfig) -> Result<RunPaths> {
-    let timestamp = Local::now().format("%Y%m%d-%H%M%S");
-    let root = run.base_dir.join(format!("{}-{timestamp}", run.name));
+    let run_id = std::env::var("RUSTRAIN_RUN_ID").ok();
+    let root = resolve_run_root(run, run_id.as_deref())?;
     let checkpoints = root.join("checkpoints");
     let logs = root.join("logs");
     let cache = root.join("cache");
@@ -1228,6 +1317,45 @@ pub fn prepare_run_directory(run: &RunConfig) -> Result<RunPaths> {
     })
 }
 
+pub fn prepare_rank_log_directory(
+    run_paths: &RunPaths,
+    rank: usize,
+    world_size: usize,
+) -> Result<PathBuf> {
+    if world_size == 0 || rank >= world_size {
+        return Err(anyhow!(
+            "invalid launcher rank {rank} for WORLD_SIZE={world_size}"
+        ));
+    }
+    let log_dir = if world_size > 1 {
+        run_paths.logs.join(format!("rank-{rank:05}"))
+    } else {
+        run_paths.logs.clone()
+    };
+    fs::create_dir_all(&log_dir)
+        .with_context(|| format!("failed to create {}", log_dir.display()))?;
+    Ok(log_dir)
+}
+
+fn resolve_run_root(run: &RunConfig, run_id: Option<&str>) -> Result<PathBuf> {
+    let suffix = match run_id {
+        Some(run_id) => {
+            if run_id.is_empty()
+                || !run_id
+                    .bytes()
+                    .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
+            {
+                return Err(anyhow!(
+                    "RUSTRAIN_RUN_ID must contain only ASCII letters, digits, '.', '_', or '-'"
+                ));
+            }
+            run_id.to_string()
+        }
+        None => Local::now().format("%Y%m%d-%H%M%S").to_string(),
+    };
+    Ok(run.base_dir.join(format!("{}-{suffix}", run.name)))
+}
+
 pub fn init_logging(log_dir: &Path) -> Result<tracing_appender::non_blocking::WorkerGuard> {
     let file_appender = tracing_appender::rolling::never(log_dir, "train.log");
     let (file_writer, guard) = tracing_appender::non_blocking(file_appender);
@@ -1300,6 +1428,64 @@ fn default_backend() -> BackendKind {
 mod tests {
     use super::*;
 
+    #[test]
+    fn shared_run_id_resolves_the_same_rank_output_directory() {
+        let run = RunConfig {
+            name: "distributed".into(),
+            base_dir: "/tmp/rustrain-tests".into(),
+            seed: 1,
+        };
+
+        let rank_zero = resolve_run_root(&run, Some("launch-123")).unwrap();
+        let rank_one = resolve_run_root(&run, Some("launch-123")).unwrap();
+
+        assert_eq!(rank_zero, rank_one);
+        assert_eq!(
+            rank_zero,
+            PathBuf::from("/tmp/rustrain-tests/distributed-launch-123")
+        );
+    }
+
+    #[test]
+    fn shared_run_id_rejects_path_components() {
+        let run = RunConfig {
+            name: "distributed".into(),
+            base_dir: "/tmp/rustrain-tests".into(),
+            seed: 1,
+        };
+
+        let error = resolve_run_root(&run, Some("../other-run")).unwrap_err();
+        assert!(error.to_string().contains("RUSTRAIN_RUN_ID"));
+    }
+
+    #[test]
+    fn distributed_rank_logs_are_isolated_under_shared_run_root() {
+        let unique = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let root = std::env::temp_dir().join(format!(
+            "rustrain-rank-logs-{}-{unique}",
+            std::process::id()
+        ));
+        std::fs::create_dir(&root).unwrap();
+        let run_paths = RunPaths {
+            root: root.clone(),
+            checkpoints: root.join("checkpoints"),
+            logs: root.join("logs"),
+            cache: root.join("cache"),
+            resolved_config: root.join("resolved_config.toml"),
+        };
+
+        let rank_zero = prepare_rank_log_directory(&run_paths, 0, 2).unwrap();
+        let rank_one = prepare_rank_log_directory(&run_paths, 1, 2).unwrap();
+
+        assert_ne!(rank_zero, rank_one);
+        assert_eq!(rank_zero, root.join("logs/rank-00000"));
+        assert_eq!(rank_one, root.join("logs/rank-00001"));
+        std::fs::remove_dir_all(root).unwrap();
+    }
+
     #[test]
     fn qwen_lora_sft_global_batch_matches_gradient_accumulation() {
         let mut config = qwen_lora_sft_config();
@@ -1316,21 +1502,85 @@ mod tests {
     #[test]
     fn cuda_memory_fraction_must_be_finite_and_bounded() {
         let mut config = qwen_lora_sft_config();
+        config.train.cuda_memory_fraction = f64::NAN;
+        assert!(validate_config(&config)
+            .expect_err("NaN CUDA memory fraction must fail")
+            .to_string()
+            .contains("cuda_memory_fraction"));
+        config.train.cuda_memory_fraction = 1.01;
+        assert!(validate_config(&config)
+            .expect_err("CUDA memory fraction above one must fail")
+            .to_string()
+            .contains("cuda_memory_fraction"));
+    }
 
-        for invalid in [0.0, -0.1, 1.01, f64::INFINITY, f64::NAN] {
-            config.train.cuda_memory_fraction = invalid;
-            let error = match validate_config(&config) {
-                Ok(()) => panic!("fraction {invalid:?} unexpectedly validated"),
-                Err(error) => error.to_string(),
-            };
-            assert!(
-                error.contains("cuda_memory_fraction must be finite and in (0, 1]"),
-                "fraction {invalid:?} unexpectedly validated or returned a different error: {error}"
-            );
-        }
+    #[test]
+    fn qwen36_lora_sft_accepts_native_projection_targets() {
+        let mut config = qwen_lora_sft_config();
+        config.model.name = "qwen3_6_test".to_string();
+        config.lora.as_mut().unwrap().target_modules = vec![
+            "q_proj".to_string(),
+            "in_proj_qkv".to_string(),
+            "in_proj_z".to_string(),
+            "out_proj".to_string(),
+        ];
+
+        config.model.architecture = "qwen3_6_lora_sft".to_string();
+        validate_config(&config).expect("native Qwen3.6 LoRA targets should validate");
+        config.model.architecture = "qwen3_5_lora_sft".to_string();
+        validate_config(&config).expect("native Qwen3.5 LoRA targets should validate");
+    }
+
+    #[test]
+    fn qwen_hybrid_lora_sft_accepts_replicated_data_parallelism() {
+        let mut config = qwen_lora_sft_config();
+        config.model.architecture = "qwen3_6_lora_sft".to_string();
+        config.parallel.data_parallel_size = 2;
+        config.train.global_batch_size = config.train.micro_batch_size
+            * config.train.gradient_accumulation_steps
+            * config.parallel.data_parallel_size;
+        validate_config(&config).expect("Qwen3.6 LoRA DP should validate");
+
+        config.model.architecture = "qwen3_5_lora_sft".to_string();
+        validate_config(&config).expect("Qwen3.5 LoRA DP should validate");
+
+        config.train.global_batch_size /= 2;
+        let error = validate_config(&config).expect_err("DP global batch mismatch should fail");
+        assert!(error.to_string().contains("data_parallel_size"));
+    }
+
+    #[test]
+    fn qwen_hybrid_lora_sft_ep_accepts_tensor_expert_data_parallelism() {
+        let mut config = qwen_lora_sft_config();
+        config.model.architecture = "qwen3_6_lora_sft_ep".to_string();
+        config.parallel.tensor_model_parallel_size = 2;
+        config.parallel.expert_model_parallel_size = 2;
+        config.parallel.data_parallel_size = 2;
+        let ep_source_sharded = std::env::var("QWEN36_EP_A2A_SHARDED")
+            .map(|value| !value.is_empty() && value != "0")
+            .unwrap_or(false);
+        config.train.global_batch_size = config.train.micro_batch_size
+            * config.train.gradient_accumulation_steps
+            * qwen_source_parallel_factor(true, true, ep_source_sharded, &config.parallel);
+        assert_eq!(
+            qwen_source_parallel_factor(true, true, true, &config.parallel),
+            4
+        );
+        assert_eq!(
+            qwen_source_parallel_factor(true, true, false, &config.parallel),
+            2
+        );
+        validate_config(&config).expect("Qwen3.6 LoRA TPxEPxDP should validate");
+
+        config.model.architecture = "qwen3_5_lora_sft_ep".to_string();
+        validate_config(&config).expect("Qwen3.5 LoRA TPxEPxDP should validate");
+
+        config.parallel.pipeline_model_parallel_size = 2;
+        validate_config(&config).expect("Qwen3.5 LoRA TPxEPxDPxPP should validate");
 
-        config.train.cuda_memory_fraction = 1.0;
-        validate_config(&config).expect("a full-device allocator fraction should validate");
+        config.parallel.context_parallel_size = 2;
+        let error = validate_config(&config).expect_err("Qwen TPxEPxDPxPPxCP should fail early");
+        assert!(error.to_string().contains("context_parallel_size"));
     }
 
     #[test]
diff --git a/crates/rustrain-ipc/src/command.rs b/crates/rustrain-ipc/src/command.rs
index 2fa3b202..8fc64cfe 100644
--- a/crates/rustrain-ipc/src/command.rs
+++ b/crates/rustrain-ipc/src/command.rs
@@ -1,10 +1,92 @@
 use serde::{Deserialize, Serialize};
 
+pub const TENSOR_SPAN_ALIGNMENT: usize = 64;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+pub struct TensorSpan {
+    pub offset_bytes: u64,
+    pub len_bytes: u64,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct TensorSlabRef {
+    pub input_ids: TensorSpan,
+    pub target_mask: TensorSpan,
+    pub attention_mask: TensorSpan,
+    pub batch_size: usize,
+    pub seq_len: usize,
+}
+
+impl TensorSlabRef {
+    pub fn spans(&self) -> [TensorSpan; 3] {
+        [self.input_ids, self.target_mask, self.attention_mask]
+    }
+
+    pub fn validate(&self, payload_len: usize) -> Result<(), String> {
+        if self.batch_size == 0 || self.seq_len == 0 {
+            return Err(format!(
+                "tensor slab batch_size and seq_len must be positive, got batch_size={} seq_len={}",
+                self.batch_size, self.seq_len
+            ));
+        }
+        let expected_bytes = self
+            .batch_size
+            .checked_mul(self.seq_len)
+            .and_then(|elements| elements.checked_mul(std::mem::size_of::<i64>()))
+            .ok_or_else(|| "tensor slab shape byte count overflowed usize".to_string())?;
+        let mut ranges = Vec::with_capacity(3);
+        for (name, span) in [
+            ("input_ids", self.input_ids),
+            ("target_mask", self.target_mask),
+            ("attention_mask", self.attention_mask),
+        ] {
+            let offset = usize::try_from(span.offset_bytes)
+                .map_err(|_| format!("{name} slab offset does not fit usize"))?;
+            let len = usize::try_from(span.len_bytes)
+                .map_err(|_| format!("{name} slab length does not fit usize"))?;
+            if offset % TENSOR_SPAN_ALIGNMENT != 0 {
+                return Err(format!(
+                    "{name} slab offset {offset} is not {TENSOR_SPAN_ALIGNMENT}-byte aligned"
+                ));
+            }
+            if len != expected_bytes {
+                return Err(format!(
+                    "{name} slab length {len} does not match batch_size={} seq_len={} ({expected_bytes} bytes)",
+                    self.batch_size, self.seq_len
+                ));
+            }
+            let end = offset
+                .checked_add(len)
+                .ok_or_else(|| format!("{name} slab range overflowed usize"))?;
+            if end > payload_len {
+                return Err(format!(
+                    "{name} slab range {offset}..{end} exceeds payload length {payload_len}"
+                ));
+            }
+            ranges.push((offset, end, name));
+        }
+        ranges.sort_unstable_by_key(|range| range.0);
+        for pair in ranges.windows(2) {
+            if pair[0].1 > pair[1].0 {
+                return Err(format!(
+                    "tensor slab spans {} and {} overlap",
+                    pair[0].2, pair[1].2
+                ));
+            }
+        }
+        Ok(())
+    }
+}
+
 /// Commands that the HTTP server can send to workers.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum EpCommand {
-    CreateSession { session_id: String },
-    DeleteSession { session_id: String },
+    CreateSession {
+        session_id: String,
+    },
+    DeleteSession {
+        session_id: String,
+    },
     LoadModel {
         session_id: String,
         model_path: String,
@@ -18,7 +100,7 @@ pub enum EpCommand {
     InitLora {
         session_id: String,
         rank: i64,
-        alpha: i64,
+        alpha: f64,
         target_layers: Vec<usize>,
         target_modules: Vec<String>,
         lr: f64,
@@ -32,6 +114,14 @@ pub enum EpCommand {
         alpha: f64,
         target_layers: Vec<i64>,
         target_modules: String,
+        #[serde(default)]
+        optimizer_lr: Option<f64>,
+        #[serde(default)]
+        optimizer_beta1: Option<f64>,
+        #[serde(default)]
+        optimizer_beta2: Option<f64>,
+        #[serde(default)]
+        optimizer_eps: Option<f64>,
     },
     BatchAddLora {
         session_id: String,
@@ -40,27 +130,59 @@ pub enum EpCommand {
         alpha: f64,
         target_layers: Vec<i64>,
         target_modules: String,
+        #[serde(default)]
+        optimizer_lr: Option<f64>,
+        #[serde(default)]
+        optimizer_beta1: Option<f64>,
+        #[serde(default)]
+        optimizer_beta2: Option<f64>,
+        #[serde(default)]
+        optimizer_eps: Option<f64>,
     },
     RemoveLora {
         session_id: String,
         adapter_id: i64,
     },
-    ListLora { session_id: String },
+    ListLora {
+        session_id: String,
+    },
     TrainStep {
         session_id: String,
         input_ids: Vec<i64>,
         target_mask: Vec<i64>,
         attention_mask: Vec<i64>,
+        #[serde(default = "default_batch_size")]
+        batch_size: usize,
         seq_len: usize,
     },
+    TrainStepSlab {
+        session_id: String,
+        tensors: TensorSlabRef,
+    },
     TrainMultiLora {
         session_id: String,
         input_ids: Vec<i64>,
         target_mask: Vec<i64>,
         attention_mask: Vec<i64>,
+        #[serde(default = "default_batch_size")]
+        batch_size: usize,
         seq_len: usize,
         n_total: i32,
         lora_rank: i32,
+        #[serde(default)]
+        adapter_ids: Vec<i64>,
+        #[serde(default)]
+        expected_steps: Vec<u64>,
+    },
+    TrainMultiLoraSlab {
+        session_id: String,
+        tensors: TensorSlabRef,
+        n_total: i32,
+        lora_rank: i32,
+        #[serde(default)]
+        adapter_ids: Vec<i64>,
+        #[serde(default)]
+        expected_steps: Vec<u64>,
     },
     EvalStep {
         session_id: String,
@@ -69,19 +191,96 @@ pub enum EpCommand {
         attention_mask: Vec<i64>,
         seq_len: usize,
     },
+    EvalStepSlab {
+        session_id: String,
+        tensors: TensorSlabRef,
+    },
+    EvalMultiLoraSlab {
+        session_id: String,
+        tensors: TensorSlabRef,
+        adapter_ids: Vec<i64>,
+    },
     ExportAdapter {
         session_id: String,
         path: String,
+        adapter_id: Option<i64>,
+        generation: String,
+    },
+    PrepareSaveCheckpoint {
+        session_id: String,
+        path: String,
+        generation: String,
+    },
+    PrepareLoadCheckpoint {
+        session_id: String,
+        path: String,
+        transaction_id: String,
+    },
+    CommitLoadCheckpoint {
+        session_id: String,
+        transaction_id: String,
+    },
+    AbortLoadCheckpoint {
+        session_id: String,
+        transaction_id: String,
+    },
+    Status {
+        session_id: String,
     },
-    Status { session_id: String },
     Shutdown,
 }
 
+impl EpCommand {
+    pub fn tensor_slab(&self) -> Option<&TensorSlabRef> {
+        match self {
+            Self::TrainStepSlab { tensors, .. }
+            | Self::TrainMultiLoraSlab { tensors, .. }
+            | Self::EvalStepSlab { tensors, .. }
+            | Self::EvalMultiLoraSlab { tensors, .. } => Some(tensors),
+            _ => None,
+        }
+    }
+}
+
+const fn default_batch_size() -> usize {
+    1
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct AdapterLoss {
+    pub adapter_id: i64,
+    pub loss: f64,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct AdapterStep {
+    pub adapter_id: i64,
+    pub step: u64,
+}
+
 /// Results that workers return to the HTTP server.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum EpResult {
     Ok,
     Loss(f64),
+    Train {
+        loss: f64,
+        step: u64,
+    },
+    MultiLoraTrain {
+        loss: f64,
+        step: u64,
+        adapter_losses: Vec<AdapterLoss>,
+        #[serde(default)]
+        adapter_steps: Vec<AdapterStep>,
+    },
+    MultiLoraEval {
+        adapter_losses: Vec<AdapterLoss>,
+    },
+    Checkpoint {
+        step: u64,
+        loss: f64,
+    },
     AdapterId(i64),
     AdapterIds(Vec<i64>),
     Count(usize),
@@ -103,3 +302,304 @@ impl EpResult {
         EpResult::Error(msg.into())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{AdapterLoss, EpCommand, EpResult, TensorSlabRef, TensorSpan};
+
+    #[test]
+    fn train_step_serde_preserves_batch_and_sequence_shape() {
+        let command = EpCommand::TrainStep {
+            session_id: "session".into(),
+            input_ids: vec![1, 2, 3, 4, 5, 6],
+            target_mask: vec![1; 6],
+            attention_mask: vec![1; 6],
+            batch_size: 2,
+            seq_len: 3,
+        };
+
+        let json = serde_json::to_string(&command).unwrap();
+        let decoded: EpCommand = serde_json::from_str(&json).unwrap();
+        match decoded {
+            EpCommand::TrainStep {
+                batch_size,
+                seq_len,
+                ..
+            } => {
+                assert_eq!(batch_size, 2);
+                assert_eq!(seq_len, 3);
+            }
+            other => panic!("unexpected command: {other:?}"),
+        }
+    }
+
+    #[test]
+    fn legacy_train_step_command_defaults_to_one_batch_row() {
+        let json = r#"{
+            "TrainStep": {
+                "session_id": "session",
+                "input_ids": [1, 2, 3],
+                "target_mask": [1, 1, 1],
+                "attention_mask": [1, 1, 1],
+                "seq_len": 3
+            }
+        }"#;
+
+        let decoded: EpCommand = serde_json::from_str(json).unwrap();
+        match decoded {
+            EpCommand::TrainStep { batch_size, .. } => assert_eq!(batch_size, 1),
+            other => panic!("unexpected command: {other:?}"),
+        }
+    }
+
+    #[test]
+    fn add_lora_optimizer_lr_preserves_default_and_explicit_semantics() {
+        let legacy = r#"{
+            "AddLora": {
+                "session_id": "session",
+                "rank": 8,
+                "alpha": 16.0,
+                "target_layers": [0],
+                "target_modules": "q_proj"
+            }
+        }"#;
+        match serde_json::from_str::<EpCommand>(legacy).unwrap() {
+            EpCommand::AddLora { optimizer_lr, .. } => assert_eq!(optimizer_lr, None),
+            other => panic!("unexpected command: {other:?}"),
+        }
+
+        let command = EpCommand::BatchAddLora {
+            session_id: "session".into(),
+            count: 3,
+            rank: 8,
+            alpha: 16.0,
+            target_layers: vec![0],
+            target_modules: "q_proj".into(),
+            optimizer_lr: Some(2.5e-4),
+            optimizer_beta1: Some(0.8),
+            optimizer_beta2: Some(0.95),
+            optimizer_eps: Some(1e-6),
+        };
+        let encoded = serde_json::to_vec(&command).unwrap();
+        match serde_json::from_slice::<EpCommand>(&encoded).unwrap() {
+            EpCommand::BatchAddLora {
+                optimizer_lr,
+                optimizer_beta1,
+                optimizer_beta2,
+                optimizer_eps,
+                ..
+            } => {
+                assert_eq!(optimizer_lr, Some(2.5e-4));
+                assert_eq!(optimizer_beta1, Some(0.8));
+                assert_eq!(optimizer_beta2, Some(0.95));
+                assert_eq!(optimizer_eps, Some(1e-6));
+            }
+            other => panic!("unexpected command: {other:?}"),
+        }
+    }
+
+    #[test]
+    fn multi_lora_result_serde_preserves_adapter_loss_order() {
+        let result = EpResult::MultiLoraTrain {
+            loss: 2.0,
+            step: 7,
+            adapter_losses: vec![
+                AdapterLoss {
+                    adapter_id: 41,
+                    loss: 1.5,
+                },
+                AdapterLoss {
+                    adapter_id: 17,
+                    loss: 2.5,
+                },
+            ],
+            adapter_steps: vec![
+                super::AdapterStep {
+                    adapter_id: 41,
+                    step: 5,
+                },
+                super::AdapterStep {
+                    adapter_id: 17,
+                    step: 9,
+                },
+            ],
+        };
+        let encoded = serde_json::to_vec(&result).unwrap();
+        let decoded: EpResult = serde_json::from_slice(&encoded).unwrap();
+        match decoded {
+            EpResult::MultiLoraTrain {
+                loss,
+                step,
+                adapter_losses,
+                adapter_steps,
+            } => {
+                assert_eq!(loss, 2.0);
+                assert_eq!(step, 7);
+                assert_eq!(adapter_losses[0].adapter_id, 41);
+                assert_eq!(adapter_losses[1].adapter_id, 17);
+                assert_eq!(adapter_steps[0].step, 5);
+                assert_eq!(adapter_steps[1].step, 9);
+            }
+            other => panic!("unexpected result: {other:?}"),
+        }
+    }
+
+    #[test]
+    fn multi_lora_expected_steps_are_backward_compatible() {
+        let legacy = r#"{
+            "TrainMultiLora": {
+                "session_id": "session",
+                "input_ids": [1, 2],
+                "target_mask": [1, 1],
+                "attention_mask": [1, 1],
+                "batch_size": 1,
+                "seq_len": 2,
+                "n_total": 2,
+                "lora_rank": 8,
+                "adapter_ids": [11, 12]
+            }
+        }"#;
+        match serde_json::from_str::<EpCommand>(legacy).unwrap() {
+            EpCommand::TrainMultiLora { expected_steps, .. } => {
+                assert!(expected_steps.is_empty());
+            }
+            other => panic!("unexpected command: {other:?}"),
+        }
+
+        let command = EpCommand::TrainMultiLora {
+            session_id: "session".into(),
+            input_ids: vec![1, 2],
+            target_mask: vec![1, 1],
+            attention_mask: vec![1, 1],
+            batch_size: 1,
+            seq_len: 2,
+            n_total: 2,
+            lora_rank: 8,
+            adapter_ids: vec![11, 12],
+            expected_steps: vec![7, 9],
+        };
+        let encoded = serde_json::to_vec(&command).unwrap();
+        match serde_json::from_slice::<EpCommand>(&encoded).unwrap() {
+            EpCommand::TrainMultiLora { expected_steps, .. } => {
+                assert_eq!(expected_steps, vec![7, 9]);
+            }
+            other => panic!("unexpected command: {other:?}"),
+        }
+    }
+
+    #[test]
+    fn train_result_serde_preserves_logical_step() {
+        let result = EpResult::Train {
+            loss: 1.25,
+            step: 7,
+        };
+        let json = serde_json::to_string(&result).unwrap();
+        let decoded: EpResult = serde_json::from_str(&json).unwrap();
+        match decoded {
+            EpResult::Train { loss, step } => {
+                assert_eq!(loss, 1.25);
+                assert_eq!(step, 7);
+            }
+            other => panic!("unexpected result: {other:?}"),
+        }
+    }
+
+    #[test]
+    fn checkpoint_commands_preserve_transaction_identity() {
+        let commands = [
+            EpCommand::PrepareSaveCheckpoint {
+                session_id: "session".into(),
+                path: "/tmp/checkpoint.partial".into(),
+                generation: "ep-1-2-3-save".into(),
+            },
+            EpCommand::PrepareLoadCheckpoint {
+                session_id: "session".into(),
+                path: "/tmp/checkpoint".into(),
+                transaction_id: "ep-1-2-4-load".into(),
+            },
+            EpCommand::CommitLoadCheckpoint {
+                session_id: "session".into(),
+                transaction_id: "ep-1-2-4-load".into(),
+            },
+            EpCommand::AbortLoadCheckpoint {
+                session_id: "session".into(),
+                transaction_id: "ep-1-2-4-load".into(),
+            },
+        ];
+
+        for command in commands {
+            let encoded = serde_json::to_string(&command).unwrap();
+            let decoded: EpCommand = serde_json::from_str(&encoded).unwrap();
+            assert_eq!(
+                serde_json::to_value(decoded).unwrap(),
+                serde_json::to_value(command).unwrap()
+            );
+        }
+    }
+
+    #[test]
+    fn checkpoint_result_serde_preserves_step_and_loss() {
+        let result = EpResult::Checkpoint {
+            step: 19,
+            loss: 0.625,
+        };
+        let encoded = serde_json::to_string(&result).unwrap();
+        let decoded: EpResult = serde_json::from_str(&encoded).unwrap();
+        match decoded {
+            EpResult::Checkpoint { step, loss } => {
+                assert_eq!(step, 19);
+                assert_eq!(loss, 0.625);
+            }
+            other => panic!("unexpected result: {other:?}"),
+        }
+    }
+
+    #[test]
+    fn tensor_slab_command_serde_preserves_spans() {
+        let tensors = TensorSlabRef {
+            input_ids: TensorSpan {
+                offset_bytes: 0,
+                len_bytes: 32,
+            },
+            target_mask: TensorSpan {
+                offset_bytes: 64,
+                len_bytes: 32,
+            },
+            attention_mask: TensorSpan {
+                offset_bytes: 128,
+                len_bytes: 32,
+            },
+            batch_size: 2,
+            seq_len: 2,
+        };
+        let command = EpCommand::TrainStepSlab {
+            session_id: "session".into(),
+            tensors: tensors.clone(),
+        };
+        let encoded = serde_json::to_vec(&command).unwrap();
+        let decoded: EpCommand = serde_json::from_slice(&encoded).unwrap();
+        assert_eq!(decoded.tensor_slab(), Some(&tensors));
+        tensors.validate(160).unwrap();
+    }
+
+    #[test]
+    fn tensor_slab_validation_rejects_overlapping_spans() {
+        let tensors = TensorSlabRef {
+            input_ids: TensorSpan {
+                offset_bytes: 0,
+                len_bytes: 16,
+            },
+            target_mask: TensorSpan {
+                offset_bytes: 0,
+                len_bytes: 16,
+            },
+            attention_mask: TensorSpan {
+                offset_bytes: 128,
+                len_bytes: 16,
+            },
+            batch_size: 1,
+            seq_len: 2,
+        };
+        assert!(tensors.validate(144).unwrap_err().contains("overlap"));
+    }
+}
diff --git a/crates/rustrain-ipc/src/lib.rs b/crates/rustrain-ipc/src/lib.rs
index f523baa1..3e564a85 100644
--- a/crates/rustrain-ipc/src/lib.rs
+++ b/crates/rustrain-ipc/src/lib.rs
@@ -1,5 +1,5 @@
-pub mod shm;
 pub mod command;
+pub mod shm;
 
-pub use command::{EpCommand, EpResult};
-pub use shm::{EpChannel, EpWorker};
+pub use command::{EpCommand, EpResult, TENSOR_SPAN_ALIGNMENT, TensorSlabRef, TensorSpan};
+pub use shm::{DEFAULT_TENSOR_SLAB_BYTES, EpChannel, EpWorker};
diff --git a/crates/rustrain-ipc/src/shm.rs b/crates/rustrain-ipc/src/shm.rs
index 9069b554..372402ac 100644
--- a/crates/rustrain-ipc/src/shm.rs
+++ b/crates/rustrain-ipc/src/shm.rs
@@ -1,24 +1,198 @@
+use std::cell::Cell;
 use std::ffi::CString;
 use std::io;
 use std::ptr;
+use std::sync::Mutex;
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::time::Duration;
 
-use crate::command::{EpCommand, EpResult};
+use crate::command::{EpCommand, EpResult, TENSOR_SPAN_ALIGNMENT, TensorSpan};
 
-/// Each slot: 4 bytes len + up to SLOT_DATA bytes JSON
-const SLOT_HEADER: usize = 4;
-const SLOT_DATA: usize = 256 * 1024; // 256KB per slot (supports seq_len up to ~32K)
+/// Each slot has a fixed transport header followed by compact JSON control data.
+const SLOT_HEADER: usize = 16;
+const SLOT_DATA: usize = 256 * 1024;
 const SLOT_SIZE: usize = SLOT_HEADER + SLOT_DATA;
+const SLOT_LEN_OFFSET: usize = 0;
+const SLOT_EPOCH_OFFSET: usize = 8;
 
 /// Shared memory layout:
 ///   [0..SLOT_SIZE)           command slot (parent writes, all workers read)
 ///   [SLOT_SIZE..SLOT_SIZE*(1+world_size))  per-worker result slots
-///
-/// Semaphores are allocated after the data region.
-const SHM_SIZE: usize = 256 * 1024;
+///   [aligned control end..semaphore end) process-shared semaphores
+///   [64-byte aligned slab header][64-byte aligned raw tensor slab]
+const SLAB_ALIGNMENT: usize = 64;
+const SLAB_HEADER_SIZE: usize = 64;
+const SLAB_MAGIC: [u8; 8] = *b"RTSLAB01";
+const SLAB_VERSION: u32 = 2;
+const HEADER_MAGIC_OFFSET: usize = 0;
+const HEADER_VERSION_OFFSET: usize = 8;
+const HEADER_WORLD_SIZE_OFFSET: usize = 12;
+const HEADER_SHM_SIZE_OFFSET: usize = 16;
+const HEADER_PAYLOAD_OFFSET_OFFSET: usize = 24;
+const HEADER_CAPACITY_OFFSET: usize = 32;
+const HEADER_EPOCH_OFFSET: usize = 40;
+const HEADER_PAYLOAD_LEN_OFFSET: usize = 48;
 
-fn cmd_offset() -> usize { 0 }
-fn result_offset(rank: usize) -> usize { SLOT_SIZE * (1 + rank) }
-fn sem_region_start(world_size: usize) -> usize { SLOT_SIZE * (1 + world_size) }
+pub const DEFAULT_TENSOR_SLAB_BYTES: usize = 32 * 1024 * 1024;
+const TENSOR_SLAB_BYTES_ENV: &str = "RUSTRAIN_EP_TENSOR_SLAB_BYTES";
+
+/// Default upper bound for one command across all workers.
+pub const DEFAULT_BROADCAST_TIMEOUT: Duration = Duration::from_secs(30 * 60);
+
+fn multi_lora_results_are_consistent(reference: &EpResult, candidate: &EpResult) -> bool {
+    match (reference, candidate) {
+        (
+            EpResult::MultiLoraTrain {
+                loss: reference_loss,
+                step: reference_step,
+                adapter_losses: reference_adapters,
+                adapter_steps: reference_steps,
+            },
+            EpResult::MultiLoraTrain {
+                loss: candidate_loss,
+                step: candidate_step,
+                adapter_losses: candidate_adapters,
+                adapter_steps: candidate_steps,
+            },
+        ) => {
+            reference_loss.to_bits() == candidate_loss.to_bits()
+                && reference_step == candidate_step
+                && reference_adapters.len() == candidate_adapters.len()
+                && reference_steps == candidate_steps
+                && reference_adapters.iter().zip(candidate_adapters).all(
+                    |(reference_adapter, candidate_adapter)| {
+                        reference_adapter.adapter_id == candidate_adapter.adapter_id
+                            && reference_adapter.loss.to_bits() == candidate_adapter.loss.to_bits()
+                    },
+                )
+        }
+        (EpResult::MultiLoraTrain { .. }, _) | (_, EpResult::MultiLoraTrain { .. }) => false,
+        _ => true,
+    }
+}
+
+fn cmd_offset() -> usize {
+    0
+}
+fn checked_align_up(value: usize, alignment: usize) -> io::Result<usize> {
+    if alignment == 0 || !alignment.is_power_of_two() {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidInput,
+            format!("alignment {alignment} is not a nonzero power of two"),
+        ));
+    }
+    value
+        .checked_add(alignment - 1)
+        .map(|aligned| aligned & !(alignment - 1))
+        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "shared layout overflow"))
+}
+
+#[derive(Debug, Clone, Copy)]
+struct ShmLayout {
+    world_size: usize,
+    sem_start: usize,
+    slab_header_offset: usize,
+    slab_payload_offset: usize,
+    slab_capacity: usize,
+    shm_size: usize,
+}
+
+impl ShmLayout {
+    fn new(world_size: usize, slab_capacity: usize) -> io::Result<Self> {
+        if world_size == 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "EP world_size must be positive",
+            ));
+        }
+        if u32::try_from(world_size).is_err() {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "EP world_size does not fit the shared layout header",
+            ));
+        }
+        let control_slots = world_size.checked_add(1).ok_or_else(layout_overflow)?;
+        let control_end = SLOT_SIZE
+            .checked_mul(control_slots)
+            .ok_or_else(layout_overflow)?;
+        let sem_alignment = std::mem::align_of::<libc::sem_t>();
+        let sem_start = checked_align_up(control_end, sem_alignment)?;
+        let sem_bytes = world_size
+            .checked_mul(2)
+            .and_then(|count| count.checked_mul(std::mem::size_of::<libc::sem_t>()))
+            .ok_or_else(layout_overflow)?;
+        let sem_end = sem_start
+            .checked_add(sem_bytes)
+            .ok_or_else(layout_overflow)?;
+        let slab_header_offset = checked_align_up(sem_end, SLAB_ALIGNMENT)?;
+        let slab_payload_offset = slab_header_offset
+            .checked_add(SLAB_HEADER_SIZE)
+            .ok_or_else(layout_overflow)?;
+        let shm_size = slab_payload_offset
+            .checked_add(slab_capacity)
+            .ok_or_else(layout_overflow)?;
+        Ok(Self {
+            world_size,
+            sem_start,
+            slab_header_offset,
+            slab_payload_offset,
+            slab_capacity,
+            shm_size,
+        })
+    }
+
+    fn result_offset(&self, rank: usize) -> io::Result<usize> {
+        if rank >= self.world_size {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!(
+                    "worker rank {rank} is outside world_size {}",
+                    self.world_size
+                ),
+            ));
+        }
+        SLOT_SIZE
+            .checked_mul(rank.checked_add(1).ok_or_else(layout_overflow)?)
+            .ok_or_else(layout_overflow)
+    }
+
+    fn semaphore_offsets(&self, rank: usize) -> io::Result<(usize, usize)> {
+        if rank >= self.world_size {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!(
+                    "worker rank {rank} is outside world_size {}",
+                    self.world_size
+                ),
+            ));
+        }
+        let sem_size = std::mem::size_of::<libc::sem_t>();
+        let request = rank
+            .checked_mul(2)
+            .and_then(|index| index.checked_mul(sem_size))
+            .and_then(|offset| self.sem_start.checked_add(offset))
+            .ok_or_else(layout_overflow)?;
+        let done = request.checked_add(sem_size).ok_or_else(layout_overflow)?;
+        Ok((request, done))
+    }
+}
+
+fn layout_overflow() -> io::Error {
+    io::Error::new(io::ErrorKind::InvalidInput, "shared memory layout overflow")
+}
+
+fn configured_slab_capacity() -> io::Result<usize> {
+    match std::env::var(TENSOR_SLAB_BYTES_ENV) {
+        Ok(value) => value.parse::<usize>().map_err(|_| {
+            io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!("{TENSOR_SLAB_BYTES_ENV} must be a non-negative integer"),
+            )
+        }),
+        Err(std::env::VarError::NotPresent) => Ok(DEFAULT_TENSOR_SLAB_BYTES),
+        Err(error) => Err(io::Error::new(io::ErrorKind::InvalidInput, error)),
+    }
+}
 
 /// Parent-side coordinator: signals all workers, waits for completion.
 pub struct EpChannel {
@@ -28,7 +202,12 @@ pub struct EpChannel {
     sem_request: Vec<*mut libc::sem_t>,
     sem_done: Vec<*mut libc::sem_t>,
     world_size: usize,
+    layout: ShmLayout,
     shm_name: String,
+    default_timeout: Duration,
+    poisoned: AtomicBool,
+    next_epoch: AtomicU64,
+    broadcast_lock: Mutex<()>,
 }
 
 /// Worker-side endpoint: waits for commands, signals completion.
@@ -38,7 +217,9 @@ pub struct EpWorker {
     sem_request: *mut libc::sem_t,
     sem_done: *mut libc::sem_t,
     rank: usize,
-    world_size: usize,
+    layout: ShmLayout,
+    current_epoch: Cell<u64>,
+    current_payload_len: Cell<usize>,
 }
 
 unsafe impl Send for EpChannel {}
@@ -47,96 +228,459 @@ unsafe impl Send for EpWorker {}
 
 impl EpChannel {
     pub fn new(world_size: usize) -> io::Result<Self> {
+        Self::new_with_timeout(world_size, DEFAULT_BROADCAST_TIMEOUT)
+    }
+
+    pub fn new_with_timeout(world_size: usize, default_timeout: Duration) -> io::Result<Self> {
+        Self::new_with_timeout_and_slab(world_size, default_timeout, configured_slab_capacity()?)
+    }
+
+    pub fn new_with_timeout_and_slab(
+        world_size: usize,
+        default_timeout: Duration,
+        slab_capacity: usize,
+    ) -> io::Result<Self> {
+        let layout = ShmLayout::new(world_size, slab_capacity)?;
+        let shm_size_i64 = i64::try_from(layout.shm_size).map_err(|_| {
+            io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "shared memory size does not fit off_t",
+            )
+        })?;
         let shm_name = format!("/rustrain-ep-{}", std::process::id());
         let c_name = CString::new(shm_name.as_str()).unwrap();
 
-        let needed = sem_region_start(world_size) + world_size * 2 * std::mem::size_of::<libc::sem_t>();
-        let shm_size = needed.max(SHM_SIZE);
-
         let fd = unsafe { libc::shm_open(c_name.as_ptr(), libc::O_CREAT | libc::O_RDWR, 0o600) };
         if fd < 0 {
             return Err(io::Error::last_os_error());
         }
-        if unsafe { libc::ftruncate(fd, shm_size as i64) } < 0 {
+        if unsafe { libc::ftruncate(fd, shm_size_i64) } < 0 {
             let e = io::Error::last_os_error();
-            unsafe { libc::close(fd); libc::shm_unlink(c_name.as_ptr()) };
+            unsafe {
+                libc::close(fd);
+                libc::shm_unlink(c_name.as_ptr())
+            };
             return Err(e);
         }
+        let allocation_error = unsafe { libc::posix_fallocate(fd, 0, shm_size_i64) };
+        if allocation_error != 0
+            && allocation_error != libc::ENOSYS
+            && allocation_error != libc::EOPNOTSUPP
+        {
+            let error = io::Error::from_raw_os_error(allocation_error);
+            unsafe {
+                libc::close(fd);
+                libc::shm_unlink(c_name.as_ptr());
+            }
+            return Err(io::Error::new(
+                error.kind(),
+                format!(
+                    "reserve {} bytes for EP shared memory: {error}",
+                    layout.shm_size
+                ),
+            ));
+        }
         let ptr = unsafe {
-            libc::mmap(ptr::null_mut(), shm_size, libc::PROT_READ | libc::PROT_WRITE, libc::MAP_SHARED, fd, 0)
+            libc::mmap(
+                ptr::null_mut(),
+                layout.shm_size,
+                libc::PROT_READ | libc::PROT_WRITE,
+                libc::MAP_SHARED,
+                fd,
+                0,
+            )
         };
         if ptr == libc::MAP_FAILED {
             let e = io::Error::last_os_error();
-            unsafe { libc::close(fd); libc::shm_unlink(c_name.as_ptr()) };
+            unsafe {
+                libc::close(fd);
+                libc::shm_unlink(c_name.as_ptr())
+            };
             return Err(e);
         }
 
-        unsafe { ptr::write_bytes(ptr as *mut u8, 0, shm_size) };
+        // The tensor slab is intentionally left untouched. Clearing it would fault in and
+        // commit the full configured capacity before the first request arrives.
+        unsafe { ptr::write_bytes(ptr as *mut u8, 0, layout.slab_payload_offset) };
+        unsafe { write_static_header(ptr as *mut u8, &layout) };
 
-        let sem_start = sem_region_start(world_size);
-        let sem_sz = std::mem::size_of::<libc::sem_t>();
         let mut sems_request = Vec::with_capacity(world_size);
         let mut sems_done = Vec::with_capacity(world_size);
 
         for i in 0..world_size {
-            let off_req = sem_start + i * 2 * sem_sz;
-            let off_done = sem_start + (i * 2 + 1) * sem_sz;
-            if unsafe { libc::sem_init((ptr as *mut u8).add(off_req) as *mut libc::sem_t, 1, 0) } != 0 {
+            let (off_req, off_done) = layout.semaphore_offsets(i)?;
+            if unsafe { libc::sem_init((ptr as *mut u8).add(off_req) as *mut libc::sem_t, 1, 0) }
+                != 0
+            {
                 return Err(io::Error::last_os_error());
             }
-            if unsafe { libc::sem_init((ptr as *mut u8).add(off_done) as *mut libc::sem_t, 1, 0) } != 0 {
+            if unsafe { libc::sem_init((ptr as *mut u8).add(off_done) as *mut libc::sem_t, 1, 0) }
+                != 0
+            {
                 return Err(io::Error::last_os_error());
             }
             sems_request.push(unsafe { (ptr as *mut u8).add(off_req) as *mut libc::sem_t });
             sems_done.push(unsafe { (ptr as *mut u8).add(off_done) as *mut libc::sem_t });
         }
 
-        tracing::info!("EP IPC: created shm '{}' ({}KB), {} workers", shm_name, shm_size / 1024, world_size);
+        tracing::info!(
+            "EP IPC: created shm '{}' ({}KB), {} workers",
+            shm_name,
+            layout.shm_size / 1024,
+            world_size
+        );
 
-        Ok(Self { shm_ptr: ptr as *mut u8, shm_size, shm_fd: fd, sem_request: sems_request, sem_done: sems_done, world_size, shm_name })
+        Ok(Self {
+            shm_ptr: ptr as *mut u8,
+            shm_size: layout.shm_size,
+            shm_fd: fd,
+            sem_request: sems_request,
+            sem_done: sems_done,
+            world_size,
+            layout,
+            shm_name,
+            default_timeout,
+            poisoned: AtomicBool::new(false),
+            next_epoch: AtomicU64::new(0),
+            broadcast_lock: Mutex::new(()),
+        })
     }
 
-    pub fn shm_name(&self) -> &str { &self.shm_name }
-    pub fn world_size(&self) -> usize { self.world_size }
+    pub fn shm_name(&self) -> &str {
+        &self.shm_name
+    }
+    pub fn world_size(&self) -> usize {
+        self.world_size
+    }
+    pub fn slab_capacity(&self) -> usize {
+        self.layout.slab_capacity
+    }
+    pub fn is_poisoned(&self) -> bool {
+        self.poisoned.load(Ordering::Acquire)
+    }
 
-    /// Send command to ALL workers, wait for ALL to complete, return rank 0's result.
+    /// Send a command to all workers and propagate any rank's error.
     pub fn broadcast(&self, cmd: &EpCommand) -> io::Result<EpResult> {
-        let json = serde_json::to_vec(cmd).map_err(|e| io::Error::other(e.to_string()))?;
+        self.broadcast_timeout(cmd, self.default_timeout)
+    }
+
+    /// Send a command with an explicit deadline suitable for bounded operations and tests.
+    pub fn broadcast_timeout(&self, cmd: &EpCommand, timeout: Duration) -> io::Result<EpResult> {
+        self.broadcast_with_slab_timeout(cmd, &[], timeout)
+    }
+
+    pub fn broadcast_with_slab(&self, cmd: &EpCommand, payload: &[u8]) -> io::Result<EpResult> {
+        self.broadcast_with_slab_timeout(cmd, payload, self.default_timeout)
+    }
+
+    pub fn broadcast_with_slab_timeout(
+        &self,
+        cmd: &EpCommand,
+        payload: &[u8],
+        timeout: Duration,
+    ) -> io::Result<EpResult> {
+        self.ensure_healthy()?;
+        let _guard = self
+            .broadcast_lock
+            .lock()
+            .map_err(|error| self.poison(format!("EP broadcast lock poisoned: {error}")))?;
+        self.ensure_healthy()?;
+
+        let json = serde_json::to_vec(cmd).map_err(|error| {
+            io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!("command serialization failed: {error}"),
+            )
+        })?;
         if json.len() > SLOT_DATA {
-            return Err(io::Error::other(format!("command too large: {} bytes", json.len())));
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!(
+                    "command too large: {} bytes exceeds {SLOT_DATA}",
+                    json.len()
+                ),
+            ));
         }
+        self.validate_slab_command(cmd, payload)?;
+        let epoch = self
+            .next_epoch
+            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
+                current.checked_add(1)
+            })
+            .map(|previous| previous + 1)
+            .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "EP IPC epoch exhausted"))?;
 
-        // Write command to shared slot (offset 0)
+        // The semaphore post publishes the immutable slab and command to every worker.
         unsafe {
-            let len = json.len() as u32;
-            ptr::copy_nonoverlapping(&len as *const u32 as *const u8, self.shm_ptr.add(cmd_offset()), SLOT_HEADER);
-            ptr::copy_nonoverlapping(json.as_ptr(), self.shm_ptr.add(cmd_offset() + SLOT_HEADER), json.len());
+            ptr::copy_nonoverlapping(
+                payload.as_ptr(),
+                self.shm_ptr.add(self.layout.slab_payload_offset),
+                payload.len(),
+            );
+            write_u64(
+                self.shm_ptr,
+                self.layout.slab_header_offset + HEADER_EPOCH_OFFSET,
+                epoch,
+            );
+            write_u64(
+                self.shm_ptr,
+                self.layout.slab_header_offset + HEADER_PAYLOAD_LEN_OFFSET,
+                payload.len() as u64,
+            );
+            write_u32(
+                self.shm_ptr,
+                cmd_offset() + SLOT_LEN_OFFSET,
+                json.len() as u32,
+            );
+            write_u64(self.shm_ptr, cmd_offset() + SLOT_EPOCH_OFFSET, epoch);
+            ptr::copy_nonoverlapping(
+                json.as_ptr(),
+                self.shm_ptr.add(cmd_offset() + SLOT_HEADER),
+                json.len(),
+            );
         }
 
         // Signal all workers simultaneously
         for i in 0..self.world_size {
             if unsafe { libc::sem_post(self.sem_request[i]) } != 0 {
-                return Err(io::Error::last_os_error());
+                let error = io::Error::last_os_error();
+                return Err(self.poison(format!("failed to signal worker rank {i}: {error}")));
             }
         }
 
-        // Wait for all workers to complete
+        // One absolute deadline bounds the complete broadcast, not each worker in turn.
+        let deadline = realtime_deadline(timeout).map_err(|error| {
+            self.poison(format!("failed to create EP broadcast deadline: {error}"))
+        })?;
         for i in 0..self.world_size {
-            if unsafe { libc::sem_wait(self.sem_done[i]) } != 0 {
-                return Err(io::Error::last_os_error());
+            if let Err(error) = sem_timedwait(self.sem_done[i], &deadline) {
+                return Err(self.poison(format!(
+                    "worker rank {i} did not complete before the EP broadcast deadline: {error}"
+                )));
+            }
+        }
+
+        let mut rank_zero = None;
+        let mut worker_error = None;
+        for rank in 0..self.world_size {
+            let result = self
+                .read_result(rank, epoch)
+                .map_err(|error| self.poison(error.to_string()))?;
+            if let EpResult::Error(error) = &result {
+                if worker_error.is_none() {
+                    worker_error = Some(EpResult::Error(format!("worker rank {rank}: {error}")));
+                }
+            }
+            if rank == 0 {
+                rank_zero = Some(result);
+            } else if worker_error.is_none()
+                && !multi_lora_results_are_consistent(
+                    rank_zero
+                        .as_ref()
+                        .expect("rank zero result must be read before later ranks"),
+                    &result,
+                )
+            {
+                worker_error = Some(EpResult::Error(format!(
+                    "worker rank {rank} returned a multi-LoRA result inconsistent with rank 0"
+                )));
             }
         }
+        if let Some(error) = worker_error {
+            return Ok(error);
+        }
+        rank_zero.ok_or_else(|| io::Error::other("EP broadcast has no rank 0 worker"))
+    }
+
+    fn validate_slab_command(&self, cmd: &EpCommand, payload: &[u8]) -> io::Result<()> {
+        if payload.len() > self.layout.slab_capacity {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!(
+                    "tensor slab payload {} bytes exceeds capacity {}",
+                    payload.len(),
+                    self.layout.slab_capacity
+                ),
+            ));
+        }
+        match cmd.tensor_slab() {
+            Some(tensors) => tensors
+                .validate(payload.len())
+                .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error)),
+            None if payload.is_empty() => Ok(()),
+            None => Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "raw tensor slab payload requires a slab command variant",
+            )),
+        }
+    }
+
+    fn ensure_healthy(&self) -> io::Result<()> {
+        if self.is_poisoned() {
+            Err(io::Error::new(
+                io::ErrorKind::BrokenPipe,
+                "EP channel is poisoned by an earlier incomplete broadcast",
+            ))
+        } else {
+            Ok(())
+        }
+    }
+
+    fn poison(&self, message: impl Into<String>) -> io::Error {
+        self.poisoned.store(true, Ordering::Release);
+        io::Error::new(io::ErrorKind::BrokenPipe, message.into())
+    }
 
-        // Read rank 0's result from its dedicated slot
-        let r0_off = result_offset(0);
+    fn read_result(&self, rank: usize, expected_epoch: u64) -> io::Result<EpResult> {
+        let offset = self.layout.result_offset(rank)?;
         unsafe {
-            let result_len = *(self.shm_ptr.add(r0_off) as *const u32) as usize;
+            let result_len = read_u32(self.shm_ptr, offset + SLOT_LEN_OFFSET) as usize;
             if result_len == 0 || result_len > SLOT_DATA {
-                return Err(io::Error::other("worker returned empty result"));
+                return Err(io::Error::other(format!(
+                    "worker rank {rank} returned an invalid result length {result_len}"
+                )));
+            }
+            let result_epoch = read_u64(self.shm_ptr, offset + SLOT_EPOCH_OFFSET);
+            if result_epoch != expected_epoch {
+                return Err(io::Error::other(format!(
+                    "worker rank {rank} returned epoch {result_epoch}, expected {expected_epoch}"
+                )));
             }
-            let result_bytes = std::slice::from_raw_parts(self.shm_ptr.add(r0_off + SLOT_HEADER), result_len);
-            serde_json::from_slice::<EpResult>(result_bytes)
-                .map_err(|e| io::Error::other(format!("result deserialization: {e}")))
+            let result_bytes =
+                std::slice::from_raw_parts(self.shm_ptr.add(offset + SLOT_HEADER), result_len);
+            serde_json::from_slice::<EpResult>(result_bytes).map_err(|error| {
+                io::Error::other(format!(
+                    "worker rank {rank} result deserialization failed: {error}"
+                ))
+            })
+        }
+    }
+}
+
+unsafe fn write_static_header(base: *mut u8, layout: &ShmLayout) {
+    unsafe {
+        ptr::copy_nonoverlapping(
+            SLAB_MAGIC.as_ptr(),
+            base.add(layout.slab_header_offset + HEADER_MAGIC_OFFSET),
+            SLAB_MAGIC.len(),
+        );
+        write_u32(
+            base,
+            layout.slab_header_offset + HEADER_VERSION_OFFSET,
+            SLAB_VERSION,
+        );
+        write_u32(
+            base,
+            layout.slab_header_offset + HEADER_WORLD_SIZE_OFFSET,
+            layout.world_size as u32,
+        );
+        write_u64(
+            base,
+            layout.slab_header_offset + HEADER_SHM_SIZE_OFFSET,
+            layout.shm_size as u64,
+        );
+        write_u64(
+            base,
+            layout.slab_header_offset + HEADER_PAYLOAD_OFFSET_OFFSET,
+            layout.slab_payload_offset as u64,
+        );
+        write_u64(
+            base,
+            layout.slab_header_offset + HEADER_CAPACITY_OFFSET,
+            layout.slab_capacity as u64,
+        );
+    }
+}
+
+unsafe fn validate_static_header(base: *const u8, layout: &ShmLayout) -> io::Result<()> {
+    let magic = unsafe {
+        std::slice::from_raw_parts(
+            base.add(layout.slab_header_offset + HEADER_MAGIC_OFFSET),
+            SLAB_MAGIC.len(),
+        )
+    };
+    if magic != SLAB_MAGIC {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            "EP tensor slab magic does not match",
+        ));
+    }
+    let version = unsafe { read_u32(base, layout.slab_header_offset + HEADER_VERSION_OFFSET) };
+    if version != SLAB_VERSION {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            format!("EP tensor slab version {version} does not match {SLAB_VERSION}"),
+        ));
+    }
+    let header_world =
+        unsafe { read_u32(base, layout.slab_header_offset + HEADER_WORLD_SIZE_OFFSET) } as usize;
+    let header_shm = unsafe { read_u64(base, layout.slab_header_offset + HEADER_SHM_SIZE_OFFSET) };
+    let header_payload = unsafe {
+        read_u64(
+            base,
+            layout.slab_header_offset + HEADER_PAYLOAD_OFFSET_OFFSET,
+        )
+    };
+    let header_capacity =
+        unsafe { read_u64(base, layout.slab_header_offset + HEADER_CAPACITY_OFFSET) };
+    if header_world != layout.world_size
+        || header_shm != layout.shm_size as u64
+        || header_payload != layout.slab_payload_offset as u64
+        || header_capacity != layout.slab_capacity as u64
+    {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            format!(
+                "EP tensor slab layout mismatch: world={header_world} shm={header_shm} payload={header_payload} capacity={header_capacity}"
+            ),
+        ));
+    }
+    Ok(())
+}
+
+unsafe fn write_u32(base: *mut u8, offset: usize, value: u32) {
+    unsafe { ptr::write_unaligned(base.add(offset).cast::<u32>(), value.to_le()) };
+}
+
+unsafe fn write_u64(base: *mut u8, offset: usize, value: u64) {
+    unsafe { ptr::write_unaligned(base.add(offset).cast::<u64>(), value.to_le()) };
+}
+
+unsafe fn read_u32(base: *const u8, offset: usize) -> u32 {
+    u32::from_le(unsafe { ptr::read_unaligned(base.add(offset).cast::<u32>()) })
+}
+
+unsafe fn read_u64(base: *const u8, offset: usize) -> u64 {
+    u64::from_le(unsafe { ptr::read_unaligned(base.add(offset).cast::<u64>()) })
+}
+
+fn realtime_deadline(timeout: Duration) -> io::Result<libc::timespec> {
+    let mut deadline = libc::timespec {
+        tv_sec: 0,
+        tv_nsec: 0,
+    };
+    if unsafe { libc::clock_gettime(libc::CLOCK_REALTIME, &mut deadline) } != 0 {
+        return Err(io::Error::last_os_error());
+    }
+
+    let timeout_secs = timeout.as_secs().min(libc::time_t::MAX as u64) as libc::time_t;
+    deadline.tv_sec = deadline.tv_sec.saturating_add(timeout_secs);
+    deadline.tv_nsec += timeout.subsec_nanos() as libc::c_long;
+    if deadline.tv_nsec >= 1_000_000_000 {
+        deadline.tv_sec = deadline.tv_sec.saturating_add(1);
+        deadline.tv_nsec -= 1_000_000_000;
+    }
+    Ok(deadline)
+}
+
+fn sem_timedwait(sem: *mut libc::sem_t, deadline: &libc::timespec) -> io::Result<()> {
+    loop {
+        if unsafe { libc::sem_timedwait(sem, deadline) } == 0 {
+            return Ok(());
+        }
+        let error = io::Error::last_os_error();
+        if error.raw_os_error() != Some(libc::EINTR) {
+            return Err(error);
         }
     }
 }
@@ -144,7 +688,10 @@ impl EpChannel {
 impl Drop for EpChannel {
     fn drop(&mut self) {
         for i in 0..self.world_size {
-            unsafe { libc::sem_destroy(self.sem_request[i]); libc::sem_destroy(self.sem_done[i]) };
+            unsafe {
+                libc::sem_destroy(self.sem_request[i]);
+                libc::sem_destroy(self.sem_done[i])
+            };
         }
         unsafe {
             libc::munmap(self.shm_ptr as *mut libc::c_void, self.shm_size);
@@ -157,67 +704,260 @@ impl Drop for EpChannel {
 
 impl EpWorker {
     pub fn attach(shm_name: &str, rank: usize, world_size: usize) -> io::Result<Self> {
+        if rank >= world_size {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!("worker rank {rank} is outside world_size {world_size}"),
+            ));
+        }
+        let base_layout = ShmLayout::new(world_size, 0)?;
         let c_name = CString::new(shm_name).unwrap();
         let fd = unsafe { libc::shm_open(c_name.as_ptr(), libc::O_RDWR, 0o600) };
-        if fd < 0 { return Err(io::Error::last_os_error()); }
+        if fd < 0 {
+            return Err(io::Error::last_os_error());
+        }
 
-        let needed = sem_region_start(world_size) + world_size * 2 * std::mem::size_of::<libc::sem_t>();
-        let shm_size = needed.max(SHM_SIZE);
+        let mut stat = std::mem::MaybeUninit::<libc::stat>::zeroed();
+        if unsafe { libc::fstat(fd, stat.as_mut_ptr()) } != 0 {
+            let error = io::Error::last_os_error();
+            unsafe { libc::close(fd) };
+            return Err(error);
+        }
+        let stat = unsafe { stat.assume_init() };
+        let shm_size = usize::try_from(stat.st_size).map_err(|_| {
+            unsafe { libc::close(fd) };
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("shared memory size {} does not fit usize", stat.st_size),
+            )
+        })?;
+        if shm_size < base_layout.slab_payload_offset {
+            unsafe { libc::close(fd) };
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!(
+                    "shared memory size {shm_size} is smaller than layout prefix {}",
+                    base_layout.slab_payload_offset
+                ),
+            ));
+        }
+        let layout = ShmLayout::new(world_size, shm_size - base_layout.slab_payload_offset)?;
 
-        let ptr = unsafe { libc::mmap(ptr::null_mut(), shm_size, libc::PROT_READ | libc::PROT_WRITE, libc::MAP_SHARED, fd, 0) };
+        let ptr = unsafe {
+            libc::mmap(
+                ptr::null_mut(),
+                shm_size,
+                libc::PROT_READ | libc::PROT_WRITE,
+                libc::MAP_SHARED,
+                fd,
+                0,
+            )
+        };
         if ptr == libc::MAP_FAILED {
             let e = io::Error::last_os_error();
             unsafe { libc::close(fd) };
             return Err(e);
         }
+        unsafe { libc::close(fd) };
 
-        let sem_start = sem_region_start(world_size);
-        let sem_sz = std::mem::size_of::<libc::sem_t>();
-        let off_req = sem_start + rank * 2 * sem_sz;
-        let off_done = sem_start + (rank * 2 + 1) * sem_sz;
+        if let Err(error) = unsafe { validate_static_header(ptr.cast(), &layout) } {
+            unsafe { libc::munmap(ptr, shm_size) };
+            return Err(error);
+        }
+        let (off_req, off_done) = layout.semaphore_offsets(rank)?;
         let sem_request = unsafe { (ptr as *mut u8).add(off_req) as *mut libc::sem_t };
         let sem_done = unsafe { (ptr as *mut u8).add(off_done) as *mut libc::sem_t };
 
         tracing::info!("EP worker {}: attached to shm '{}'", rank, shm_name);
-        Ok(Self { shm_ptr: ptr as *mut u8, shm_size, sem_request, sem_done, rank, world_size })
+        Ok(Self {
+            shm_ptr: ptr as *mut u8,
+            shm_size,
+            sem_request,
+            sem_done,
+            rank,
+            layout,
+            current_epoch: Cell::new(0),
+            current_payload_len: Cell::new(0),
+        })
     }
 
     pub fn wait_command(&self) -> io::Result<EpCommand> {
-        if unsafe { libc::sem_wait(self.sem_request) } != 0 {
-            return Err(io::Error::last_os_error());
+        loop {
+            if unsafe { libc::sem_wait(self.sem_request) } == 0 {
+                break;
+            }
+            let error = io::Error::last_os_error();
+            if error.raw_os_error() != Some(libc::EINTR) {
+                return Err(error);
+            }
         }
         unsafe {
-            let cmd_len = *(self.shm_ptr.add(cmd_offset()) as *const u32) as usize;
+            let cmd_len = read_u32(self.shm_ptr, cmd_offset() + SLOT_LEN_OFFSET) as usize;
             if cmd_len == 0 || cmd_len > SLOT_DATA {
-                return Err(io::Error::other(format!("invalid command length: {}", cmd_len)));
+                return Err(io::Error::other(format!(
+                    "invalid command length: {}",
+                    cmd_len
+                )));
+            }
+            let command_epoch = read_u64(self.shm_ptr, cmd_offset() + SLOT_EPOCH_OFFSET);
+            let slab_epoch = read_u64(
+                self.shm_ptr,
+                self.layout.slab_header_offset + HEADER_EPOCH_OFFSET,
+            );
+            if command_epoch == 0 || command_epoch != slab_epoch {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    format!("command epoch {command_epoch} does not match slab epoch {slab_epoch}"),
+                ));
+            }
+            if command_epoch <= self.current_epoch.get() {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    format!(
+                        "command epoch {command_epoch} is not newer than worker epoch {}",
+                        self.current_epoch.get()
+                    ),
+                ));
             }
-            let cmd_bytes = std::slice::from_raw_parts(self.shm_ptr.add(cmd_offset() + SLOT_HEADER), cmd_len);
-            serde_json::from_slice::<EpCommand>(cmd_bytes)
-                .map_err(|e| io::Error::other(format!("command deserialization: {e}")))
+            let payload_len_u64 = read_u64(
+                self.shm_ptr,
+                self.layout.slab_header_offset + HEADER_PAYLOAD_LEN_OFFSET,
+            );
+            let payload_len = usize::try_from(payload_len_u64).map_err(|_| {
+                io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    "tensor slab payload length does not fit usize",
+                )
+            })?;
+            if payload_len > self.layout.slab_capacity {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    format!(
+                        "tensor slab payload {payload_len} exceeds mapped capacity {}",
+                        self.layout.slab_capacity
+                    ),
+                ));
+            }
+            let cmd_bytes =
+                std::slice::from_raw_parts(self.shm_ptr.add(cmd_offset() + SLOT_HEADER), cmd_len);
+            let command = serde_json::from_slice::<EpCommand>(cmd_bytes)
+                .map_err(|e| io::Error::other(format!("command deserialization: {e}")))?;
+            match command.tensor_slab() {
+                Some(tensors) => tensors
+                    .validate(payload_len)
+                    .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?,
+                None if payload_len == 0 => {}
+                None => {
+                    return Err(io::Error::new(
+                        io::ErrorKind::InvalidData,
+                        "legacy command unexpectedly carries a tensor slab payload",
+                    ));
+                }
+            }
+            self.current_epoch.set(command_epoch);
+            self.current_payload_len.set(payload_len);
+            Ok(command)
         }
     }
 
-    pub fn signal_done(&self, result: &EpResult) -> io::Result<()> {
-        let json = serde_json::to_vec(result).map_err(|e| io::Error::other(e.to_string()))?;
-        if json.len() > SLOT_DATA {
-            return Err(io::Error::other(format!("result too large: {} bytes", json.len())));
+    /// Returns a zero-copy view into the current broadcast's raw little-endian i64 slab.
+    ///
+    /// # Safety
+    /// The view must not be retained or accessed after `signal_done`, because the parent may
+    /// reuse the shared slab as soon as every worker reports completion.
+    pub unsafe fn slab_i64(&self, span: TensorSpan) -> io::Result<&[i64]> {
+        if cfg!(target_endian = "big") {
+            return Err(io::Error::new(
+                io::ErrorKind::Unsupported,
+                "tensor slab i64 views require a little-endian host",
+            ));
+        }
+        let offset = usize::try_from(span.offset_bytes).map_err(|_| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                "tensor slab offset does not fit usize",
+            )
+        })?;
+        let len = usize::try_from(span.len_bytes).map_err(|_| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                "tensor slab length does not fit usize",
+            )
+        })?;
+        if offset % TENSOR_SPAN_ALIGNMENT != 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("tensor slab offset {offset} is not {TENSOR_SPAN_ALIGNMENT}-byte aligned"),
+            ));
+        }
+        if len % std::mem::size_of::<i64>() != 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("tensor slab length {len} is not a multiple of 8 bytes"),
+            ));
+        }
+        let end = offset.checked_add(len).ok_or_else(|| {
+            io::Error::new(io::ErrorKind::InvalidData, "tensor slab span overflow")
+        })?;
+        if end > self.current_payload_len.get() {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!(
+                    "tensor slab span {offset}..{end} exceeds current payload {}",
+                    self.current_payload_len.get()
+                ),
+            ));
+        }
+        let pointer = unsafe { self.shm_ptr.add(self.layout.slab_payload_offset + offset) };
+        if pointer.align_offset(std::mem::align_of::<i64>()) != 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                "mapped tensor slab pointer is not i64 aligned",
+            ));
         }
+        Ok(unsafe { std::slice::from_raw_parts(pointer.cast::<i64>(), len / 8) })
+    }
+
+    pub fn signal_done(&self, result: &EpResult) -> io::Result<()> {
+        let json = encode_result_for_slot(result)?;
 
         // Write result to THIS worker's dedicated slot (no collision with other workers)
-        let off = result_offset(self.rank);
+        let off = self.layout.result_offset(self.rank)?;
         unsafe {
-            let len = json.len() as u32;
-            ptr::copy_nonoverlapping(&len as *const u32 as *const u8, self.shm_ptr.add(off), SLOT_HEADER);
-            ptr::copy_nonoverlapping(json.as_ptr(), self.shm_ptr.add(off + SLOT_HEADER), json.len());
+            write_u32(self.shm_ptr, off + SLOT_LEN_OFFSET, json.len() as u32);
+            write_u64(
+                self.shm_ptr,
+                off + SLOT_EPOCH_OFFSET,
+                self.current_epoch.get(),
+            );
+            ptr::copy_nonoverlapping(
+                json.as_ptr(),
+                self.shm_ptr.add(off + SLOT_HEADER),
+                json.len(),
+            );
         }
 
         if unsafe { libc::sem_post(self.sem_done) } != 0 {
             return Err(io::Error::last_os_error());
         }
+        self.current_payload_len.set(0);
         Ok(())
     }
 
-    pub fn rank(&self) -> usize { self.rank }
+    pub fn rank(&self) -> usize {
+        self.rank
+    }
+}
+
+fn encode_result_for_slot(result: &EpResult) -> io::Result<Vec<u8>> {
+    let json = serde_json::to_vec(result).map_err(|e| io::Error::other(e.to_string()))?;
+    if json.len() <= SLOT_DATA {
+        return Ok(json);
+    }
+    serde_json::to_vec(&EpResult::Error(format!(
+        "worker result exceeded the {SLOT_DATA}-byte IPC slot"
+    )))
+    .map_err(|e| io::Error::other(e.to_string()))
 }
 
 impl Drop for EpWorker {
@@ -230,10 +970,291 @@ impl Drop for EpWorker {
 mod tests {
     use super::*;
 
+    #[test]
+    fn multi_lora_result_consistency_checks_rank_values_and_shape() {
+        let reference = EpResult::MultiLoraTrain {
+            loss: 2.0,
+            step: 7,
+            adapter_losses: vec![crate::command::AdapterLoss {
+                adapter_id: 41,
+                loss: 2.0,
+            }],
+            adapter_steps: vec![crate::command::AdapterStep {
+                adapter_id: 41,
+                step: 3,
+            }],
+        };
+        assert!(multi_lora_results_are_consistent(&reference, &reference));
+
+        let different_loss = EpResult::MultiLoraTrain {
+            loss: 2.0,
+            step: 7,
+            adapter_losses: vec![crate::command::AdapterLoss {
+                adapter_id: 41,
+                loss: 2.5,
+            }],
+            adapter_steps: vec![crate::command::AdapterStep {
+                adapter_id: 41,
+                step: 3,
+            }],
+        };
+        assert!(!multi_lora_results_are_consistent(
+            &reference,
+            &different_loss
+        ));
+        assert!(!multi_lora_results_are_consistent(
+            &reference,
+            &EpResult::Train { loss: 2.0, step: 7 }
+        ));
+    }
+
+    #[test]
+    fn oversized_result_is_replaced_by_a_compact_error() {
+        let encoded = encode_result_for_slot(&EpResult::Error("x".repeat(SLOT_DATA))).unwrap();
+        assert!(encoded.len() <= SLOT_DATA);
+        let decoded: EpResult = serde_json::from_slice(&encoded).unwrap();
+        match decoded {
+            EpResult::Error(message) => assert!(message.contains("exceeded")),
+            other => panic!("unexpected oversized-result replacement: {other:?}"),
+        }
+    }
+    use crate::command::TensorSlabRef;
+
+    static TEST_CHANNEL_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
+
+    fn test_channel(world_size: usize, slab_capacity: usize) -> EpChannel {
+        EpChannel::new_with_timeout_and_slab(world_size, DEFAULT_BROADCAST_TIMEOUT, slab_capacity)
+            .expect("create test channel")
+    }
+
+    fn append_i64(payload: &mut Vec<u8>, values: &[i64]) -> TensorSpan {
+        let aligned = payload
+            .len()
+            .checked_add(crate::command::TENSOR_SPAN_ALIGNMENT - 1)
+            .map(|value| value & !(crate::command::TENSOR_SPAN_ALIGNMENT - 1))
+            .unwrap();
+        payload.resize(aligned, 0);
+        let offset_bytes = payload.len() as u64;
+        for value in values {
+            payload.extend_from_slice(&value.to_le_bytes());
+        }
+        TensorSpan {
+            offset_bytes,
+            len_bytes: (values.len() * std::mem::size_of::<i64>()) as u64,
+        }
+    }
+
     #[test]
     fn test_channel_create_destroy() {
-        let ch = EpChannel::new(2).expect("create channel");
+        let _guard = TEST_CHANNEL_LOCK.lock().unwrap();
+        let ch = test_channel(2, 4096);
         assert_eq!(ch.world_size(), 2);
+        assert_eq!(ch.slab_capacity(), 4096);
         drop(ch);
     }
+
+    #[test]
+    fn checked_layout_aligns_semaphores_and_tensor_slab() {
+        for world_size in [1, 2, 8] {
+            let layout = ShmLayout::new(world_size, 4096).unwrap();
+            assert_eq!(layout.sem_start % std::mem::align_of::<libc::sem_t>(), 0);
+            assert_eq!(layout.slab_header_offset % SLAB_ALIGNMENT, 0);
+            assert_eq!(layout.slab_payload_offset % SLAB_ALIGNMENT, 0);
+            assert_eq!(layout.shm_size - layout.slab_payload_offset, 4096);
+        }
+        assert_eq!(
+            ShmLayout::new(0, 4096).unwrap_err().kind(),
+            io::ErrorKind::InvalidInput
+        );
+    }
+
+    #[test]
+    fn broadcast_propagates_nonzero_rank_error() {
+        let _guard = TEST_CHANNEL_LOCK.lock().unwrap();
+        let channel = test_channel(2, 4096);
+        let worker_zero = EpWorker::attach(channel.shm_name(), 0, 2).expect("attach rank 0");
+        let worker_one = EpWorker::attach(channel.shm_name(), 1, 2).expect("attach rank 1");
+        let rank_zero = std::thread::spawn(move || {
+            assert!(matches!(
+                worker_zero.wait_command().unwrap(),
+                EpCommand::Shutdown
+            ));
+            worker_zero.signal_done(&EpResult::Ok).unwrap();
+        });
+        let rank_one = std::thread::spawn(move || {
+            assert!(matches!(
+                worker_one.wait_command().unwrap(),
+                EpCommand::Shutdown
+            ));
+            worker_one
+                .signal_done(&EpResult::Error("rank one failed".into()))
+                .unwrap();
+        });
+
+        let result = channel.broadcast(&EpCommand::Shutdown).unwrap();
+        rank_zero.join().unwrap();
+        rank_one.join().unwrap();
+        assert!(!channel.is_poisoned());
+
+        match result {
+            EpResult::Error(error) => {
+                assert_eq!(error, "worker rank 1: rank one failed");
+            }
+            _ => panic!("rank 1 error was not propagated"),
+        }
+    }
+
+    #[test]
+    fn broadcast_timeout_permanently_poisons_channel() {
+        let _guard = TEST_CHANNEL_LOCK.lock().unwrap();
+        let channel = test_channel(1, 4096);
+        let worker = EpWorker::attach(channel.shm_name(), 0, 1).expect("attach rank 0");
+        let worker_thread = std::thread::spawn(move || {
+            assert!(matches!(
+                worker.wait_command().unwrap(),
+                EpCommand::Shutdown
+            ));
+            // Simulate a worker crash after accepting the command.
+        });
+
+        let started = std::time::Instant::now();
+        let error = channel
+            .broadcast_timeout(&EpCommand::Shutdown, Duration::from_millis(100))
+            .expect_err("missing completion must time out");
+        assert_eq!(error.kind(), io::ErrorKind::BrokenPipe);
+        assert!(started.elapsed() < Duration::from_secs(2));
+        assert!(channel.is_poisoned());
+        worker_thread.join().unwrap();
+
+        let retry_started = std::time::Instant::now();
+        let retry_error = channel
+            .broadcast_timeout(&EpCommand::Shutdown, Duration::from_secs(1))
+            .expect_err("poisoned channel must reject future commands");
+        assert_eq!(retry_error.kind(), io::ErrorKind::BrokenPipe);
+        assert!(retry_started.elapsed() < Duration::from_millis(100));
+    }
+
+    #[test]
+    fn tensor_slab_round_trip_exposes_zero_copy_i64_views() {
+        let _guard = TEST_CHANNEL_LOCK.lock().unwrap();
+        let channel = test_channel(1, 256);
+        let worker = EpWorker::attach(channel.shm_name(), 0, 1).expect("attach rank 0");
+        let mut payload = Vec::new();
+        let input_ids = append_i64(&mut payload, &[11, 12, 13, 14]);
+        let target_mask = append_i64(&mut payload, &[1, 0, 1, 0]);
+        let attention_mask = append_i64(&mut payload, &[1, 1, 1, 1]);
+        let tensors = TensorSlabRef {
+            input_ids,
+            target_mask,
+            attention_mask,
+            batch_size: 2,
+            seq_len: 2,
+        };
+        let worker_thread = std::thread::spawn(move || {
+            let command = worker.wait_command().unwrap();
+            let tensors = match command {
+                EpCommand::TrainStepSlab { tensors, .. } => tensors,
+                other => panic!("unexpected command: {other:?}"),
+            };
+            unsafe {
+                assert_eq!(
+                    worker.slab_i64(tensors.input_ids).unwrap(),
+                    [11, 12, 13, 14]
+                );
+                assert_eq!(worker.slab_i64(tensors.target_mask).unwrap(), [1, 0, 1, 0]);
+                assert_eq!(
+                    worker.slab_i64(tensors.attention_mask).unwrap(),
+                    [1, 1, 1, 1]
+                );
+            }
+            worker.signal_done(&EpResult::Ok).unwrap();
+        });
+        let result = channel
+            .broadcast_with_slab(
+                &EpCommand::TrainStepSlab {
+                    session_id: "session".into(),
+                    tensors,
+                },
+                &payload,
+            )
+            .unwrap();
+        assert!(matches!(result, EpResult::Ok));
+        assert!(!channel.is_poisoned());
+        worker_thread.join().unwrap();
+    }
+
+    #[test]
+    fn invalid_or_oversized_slab_is_recoverable_before_publish() {
+        let _guard = TEST_CHANNEL_LOCK.lock().unwrap();
+        let channel = test_channel(1, 32);
+        let oversized = TensorSlabRef {
+            input_ids: TensorSpan {
+                offset_bytes: 0,
+                len_bytes: 16,
+            },
+            target_mask: TensorSpan {
+                offset_bytes: 64,
+                len_bytes: 16,
+            },
+            attention_mask: TensorSpan {
+                offset_bytes: 128,
+                len_bytes: 16,
+            },
+            batch_size: 1,
+            seq_len: 2,
+        };
+        let error = channel
+            .broadcast_with_slab(
+                &EpCommand::TrainStepSlab {
+                    session_id: "session".into(),
+                    tensors: oversized,
+                },
+                &[0; 144],
+            )
+            .unwrap_err();
+        assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
+        assert!(!channel.is_poisoned());
+
+        let overlapping = TensorSlabRef {
+            input_ids: TensorSpan {
+                offset_bytes: 0,
+                len_bytes: 8,
+            },
+            target_mask: TensorSpan {
+                offset_bytes: 0,
+                len_bytes: 8,
+            },
+            attention_mask: TensorSpan {
+                offset_bytes: 64,
+                len_bytes: 8,
+            },
+            batch_size: 1,
+            seq_len: 1,
+        };
+        let error = channel
+            .broadcast_with_slab(
+                &EpCommand::TrainStepSlab {
+                    session_id: "session".into(),
+                    tensors: overlapping,
+                },
+                &[0; 72],
+            )
+            .unwrap_err();
+        assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
+        assert!(!channel.is_poisoned());
+
+        let worker = EpWorker::attach(channel.shm_name(), 0, 1).expect("attach rank 0");
+        let worker_thread = std::thread::spawn(move || {
+            assert!(matches!(
+                worker.wait_command().unwrap(),
+                EpCommand::Shutdown
+            ));
+            worker.signal_done(&EpResult::Ok).unwrap();
+        });
+        assert!(matches!(
+            channel.broadcast(&EpCommand::Shutdown).unwrap(),
+            EpResult::Ok
+        ));
+        worker_thread.join().unwrap();
+    }
 }
diff --git a/crates/rustrain-parallel/src/launcher.rs b/crates/rustrain-parallel/src/launcher.rs
index 311002bd..b033c50d 100644
--- a/crates/rustrain-parallel/src/launcher.rs
+++ b/crates/rustrain-parallel/src/launcher.rs
@@ -9,6 +9,8 @@ use std::{
 use anyhow::{Context, Result, bail};
 use serde::{Deserialize, Serialize};
 
+use crate::topology::ParallelTopology;
+
 #[derive(Debug, Serialize)]
 struct LaunchSummary {
     nproc_per_node: usize,
@@ -29,6 +31,17 @@ struct RankSummary {
     log_path: String,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+struct LaunchNodeMarker {
+    nnodes: usize,
+    nproc_per_node: usize,
+    node_rank: usize,
+    master_addr: String,
+    master_port: u16,
+    run_id: String,
+    attempt_id: String,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct LaunchEnvSummary {
     pub rank: usize,
@@ -40,6 +53,12 @@ pub struct LaunchEnvSummary {
     pub cuda_visible_devices: Option<String>,
     pub assigned_cuda_visible_device: Option<String>,
     pub assigned_cuda_device_ordinal: Option<usize>,
+    pub tensor_model_parallel_size: usize,
+    pub pipeline_model_parallel_size: usize,
+    pub data_parallel_size: usize,
+    pub expert_model_parallel_size: usize,
+    pub context_parallel_size: usize,
+    pub parallel_rank_order: String,
 }
 
 pub fn launch(
@@ -49,7 +68,15 @@ pub fn launch(
     master_port: u16,
     command: &[String],
 ) -> Result<()> {
-    launch_multi(nproc_per_node, 1, 0, output_dir, master_addr, master_port, command)
+    launch_multi(
+        nproc_per_node,
+        1,
+        0,
+        output_dir,
+        master_addr,
+        master_port,
+        command,
+    )
 }
 
 pub fn launch_multi(
@@ -81,9 +108,25 @@ pub fn launch_multi(
         .with_context(|| format!("failed to create {}", output_dir.display()))?;
     let current_exe = std::env::current_exe().context("failed to locate current executable")?;
     let timeout = launch_timeout()?;
+    let topology = ParallelTopology::from_env_with_world_size(world_size)?;
+    let run_id = resolve_launch_run_id(nnodes, std::env::var("RUSTRAIN_RUN_ID").ok())?;
+    let attempt_id = resolve_launch_attempt_id(nnodes, std::env::var("RUSTRAIN_ATTEMPT_ID").ok())?;
     let visible_cuda_devices =
         parse_visible_cuda_devices(std::env::var("CUDA_VISIBLE_DEVICES").ok());
     validate_visible_cuda_devices(nproc_per_node, visible_cuda_devices.as_deref())?;
+    if nnodes > 1 {
+        rendezvous_launch_nodes(
+            output_dir,
+            nnodes,
+            nproc_per_node,
+            node_rank,
+            master_addr,
+            master_port,
+            &run_id,
+            &attempt_id,
+            timeout.unwrap_or(Duration::from_secs(120)),
+        )?;
+    }
 
     let mut children = Vec::with_capacity(nproc_per_node);
     for local_rank in 0..nproc_per_node {
@@ -105,8 +148,51 @@ pub fn launch_multi(
             .env("MASTER_ADDR", master_addr)
             .env("MASTER_PORT", master_port.to_string())
             .env("RUSTRAIN_LAUNCH_OUTPUT_DIR", output_dir)
+            .env("RUSTRAIN_RUN_ID", &run_id)
+            .env("RUSTRAIN_ATTEMPT_ID", &attempt_id)
             .env("NNODES", nnodes.to_string())
             .env("NODE_RANK", node_rank.to_string())
+            // Normalize topology variables for every child. Explicit
+            // TP_SIZE/PP_SIZE/... values come from the parent environment;
+            // unspecified axes default to replicated DP and are validated
+            // against the launcher's world size above.
+            .env("TP_SIZE", topology.tensor_model_parallel_size().to_string())
+            .env(
+                "PP_SIZE",
+                topology.pipeline_model_parallel_size().to_string(),
+            )
+            .env("DP_SIZE", topology.data_parallel_size().to_string())
+            .env("EP_SIZE", topology.expert_model_parallel_size().to_string())
+            .env("CP_SIZE", topology.context_parallel_size().to_string())
+            .env(
+                "RUSTRAIN_TP_SIZE",
+                topology.tensor_model_parallel_size().to_string(),
+            )
+            .env(
+                "RUSTRAIN_PP_SIZE",
+                topology.pipeline_model_parallel_size().to_string(),
+            )
+            .env(
+                "RUSTRAIN_DP_SIZE",
+                topology.data_parallel_size().to_string(),
+            )
+            .env(
+                "RUSTRAIN_EP_SIZE",
+                topology.expert_model_parallel_size().to_string(),
+            )
+            .env(
+                "RUSTRAIN_CP_SIZE",
+                topology.context_parallel_size().to_string(),
+            )
+            .env(
+                "RUSTRAIN_PARALLEL_ORDER",
+                topology
+                    .order()
+                    .into_iter()
+                    .map(|axis| axis.name())
+                    .collect::<Vec<_>>()
+                    .join("-"),
+            )
             .env(
                 "PYTORCH_CUDA_ALLOC_CONF",
                 "expandable_segments:True,max_split_size_mb:512",
@@ -119,13 +205,19 @@ pub fn launch_multi(
         {
             child
                 .env("RUSTRAIN_ASSIGNED_CUDA_VISIBLE_DEVICE", assigned_device)
-                .env("RUSTRAIN_ASSIGNED_CUDA_DEVICE_ORDINAL", local_rank.to_string());
+                .env(
+                    "RUSTRAIN_ASSIGNED_CUDA_DEVICE_ORDINAL",
+                    local_rank.to_string(),
+                );
         }
         children.push((
             global_rank,
             log_path,
             child.spawn().with_context(|| {
-                format!("failed to spawn rank {global_rank} for command {:?}", command)
+                format!(
+                    "failed to spawn rank {global_rank} for command {:?}",
+                    command
+                )
             })?,
         ));
     }
@@ -134,21 +226,22 @@ pub fn launch_multi(
     let mut ranks = Vec::with_capacity(nproc_per_node);
     let mut failed = Vec::new();
     for wait_result in wait_results {
+        let local_rank = wait_result.rank - node_rank * nproc_per_node;
         if !wait_result.success() {
             failed.push(wait_result.rank);
         }
         ranks.push(RankSummary {
             rank: wait_result.rank,
-            local_rank: wait_result.rank,
-            world_size: nproc_per_node,
+            local_rank,
+            world_size,
             assigned_cuda_visible_device: visible_cuda_devices
                 .as_ref()
-                .and_then(|devices| devices.get(wait_result.rank))
+                .and_then(|devices| devices.get(local_rank))
                 .cloned(),
             assigned_cuda_device_ordinal: visible_cuda_devices
                 .as_ref()
-                .and_then(|devices| devices.get(wait_result.rank))
-                .map(|_| wait_result.rank),
+                .and_then(|devices| devices.get(local_rank))
+                .map(|_| local_rank),
             status_code: wait_result.status.and_then(|status| status.code()),
             timed_out: wait_result.timed_out,
             log_path: wait_result.log_path.display().to_string(),
@@ -162,7 +255,11 @@ pub fn launch_multi(
         ranks,
     };
     let summary_json = serde_json::to_string_pretty(&summary)?;
-    let summary_path = output_dir.join("launch-summary.json");
+    let summary_path = if nnodes > 1 {
+        output_dir.join(format!("launch-summary-node-{node_rank}.json"))
+    } else {
+        output_dir.join("launch-summary.json")
+    };
     fs::write(&summary_path, &summary_json)
         .with_context(|| format!("failed to write {}", summary_path.display()))?;
     println!("{summary_json}");
@@ -181,6 +278,131 @@ pub fn launch_multi(
     Ok(())
 }
 
+fn resolve_launch_run_id(nnodes: usize, configured: Option<String>) -> Result<String> {
+    if let Some(configured) = configured.filter(|value| !value.is_empty()) {
+        validate_launch_id("RUSTRAIN_RUN_ID", &configured)?;
+        return Ok(configured);
+    }
+    if nnodes > 1 {
+        bail!("multi-node launch requires one shared RUSTRAIN_RUN_ID to be set on every launcher");
+    }
+    let nonce = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .context("system clock is before the Unix epoch")?
+        .as_nanos();
+    Ok(format!("launch-{}-{nonce}", std::process::id()))
+}
+
+fn resolve_launch_attempt_id(nnodes: usize, configured: Option<String>) -> Result<String> {
+    if let Some(configured) = configured.filter(|value| !value.is_empty()) {
+        validate_launch_id("RUSTRAIN_ATTEMPT_ID", &configured)?;
+        return Ok(configured);
+    }
+    if nnodes > 1 {
+        bail!(
+            "multi-node launch requires one shared RUSTRAIN_ATTEMPT_ID to be set on every launcher"
+        );
+    }
+    let nonce = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .context("system clock is before the Unix epoch")?
+        .as_nanos();
+    Ok(format!("attempt-{}-{nonce}", std::process::id()))
+}
+
+fn validate_launch_id(name: &str, value: &str) -> Result<()> {
+    if value.is_empty()
+        || matches!(value, "." | "..")
+        || !value
+            .bytes()
+            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
+    {
+        bail!(
+            "{name} may contain only ASCII letters, digits, '-', '_', and '.', and cannot be '.' or '..'"
+        );
+    }
+    Ok(())
+}
+
+#[allow(clippy::too_many_arguments)]
+fn rendezvous_launch_nodes(
+    output_dir: &Path,
+    nnodes: usize,
+    nproc_per_node: usize,
+    node_rank: usize,
+    master_addr: &str,
+    master_port: u16,
+    run_id: &str,
+    attempt_id: &str,
+    timeout: Duration,
+) -> Result<()> {
+    let rendezvous = output_dir
+        .join(".rustrain-launch")
+        .join(run_id)
+        .join(attempt_id);
+    fs::create_dir_all(&rendezvous)
+        .with_context(|| format!("failed to create {}", rendezvous.display()))?;
+    let marker = LaunchNodeMarker {
+        nnodes,
+        nproc_per_node,
+        node_rank,
+        master_addr: master_addr.to_string(),
+        master_port,
+        run_id: run_id.to_string(),
+        attempt_id: attempt_id.to_string(),
+    };
+    let marker_path = rendezvous.join(format!("node-{node_rank:05}.json"));
+    let partial_path = rendezvous.join(format!(
+        ".node-{node_rank:05}-{}.partial",
+        std::process::id()
+    ));
+    fs::write(&partial_path, serde_json::to_vec_pretty(&marker)?)
+        .with_context(|| format!("failed to write {}", partial_path.display()))?;
+    fs::rename(&partial_path, &marker_path).with_context(|| {
+        format!(
+            "failed to publish launch rendezvous marker {}",
+            marker_path.display()
+        )
+    })?;
+
+    let deadline = Instant::now() + timeout;
+    loop {
+        let mut ready = true;
+        for expected_rank in 0..nnodes {
+            let path = rendezvous.join(format!("node-{expected_rank:05}.json"));
+            let Some(contents) = fs::read(&path).ok() else {
+                ready = false;
+                break;
+            };
+            let observed: LaunchNodeMarker = serde_json::from_slice(&contents)
+                .with_context(|| format!("failed to parse {}", path.display()))?;
+            if observed.nnodes != nnodes
+                || observed.nproc_per_node != nproc_per_node
+                || observed.node_rank != expected_rank
+                || observed.master_addr != master_addr
+                || observed.master_port != master_port
+                || observed.run_id != run_id
+                || observed.attempt_id != attempt_id
+            {
+                bail!(
+                    "multi-node launch rendezvous metadata differs at {}",
+                    path.display()
+                );
+            }
+        }
+        if ready {
+            return Ok(());
+        }
+        if Instant::now() >= deadline {
+            bail!(
+                "timed out waiting for all launch nodes under {}; multi-node output_dir must be on shared storage",
+                rendezvous.display()
+            );
+        }
+        sleep(Duration::from_millis(100));
+    }
+}
+
 pub fn print_launch_env() -> Result<()> {
     let summary = read_launch_env()?;
     println!("{}", serde_json::to_string_pretty(&summary)?);
@@ -188,6 +410,7 @@ pub fn print_launch_env() -> Result<()> {
 }
 
 fn read_launch_env() -> Result<LaunchEnvSummary> {
+    let topology = ParallelTopology::from_env()?;
     Ok(LaunchEnvSummary {
         rank: parse_env_usize("RANK")?,
         local_rank: parse_env_usize("LOCAL_RANK")?,
@@ -204,6 +427,17 @@ fn read_launch_env() -> Result<LaunchEnvSummary> {
                     .with_context(|| "RUSTRAIN_ASSIGNED_CUDA_DEVICE_ORDINAL must be a usize")
             })
             .transpose()?,
+        tensor_model_parallel_size: topology.tensor_model_parallel_size(),
+        pipeline_model_parallel_size: topology.pipeline_model_parallel_size(),
+        data_parallel_size: topology.data_parallel_size(),
+        expert_model_parallel_size: topology.expert_model_parallel_size(),
+        context_parallel_size: topology.context_parallel_size(),
+        parallel_rank_order: topology
+            .order()
+            .into_iter()
+            .map(|axis| axis.name())
+            .collect::<Vec<_>>()
+            .join("-"),
     })
 }
 
@@ -337,6 +571,86 @@ mod tests {
         assert!(error.to_string().contains("requires a child command"));
     }
 
+    #[test]
+    fn launch_reuses_explicit_run_id_across_nodes() {
+        assert_eq!(
+            resolve_launch_run_id(2, Some("shared-run".into())).unwrap(),
+            "shared-run"
+        );
+    }
+
+    #[test]
+    fn multi_node_launch_requires_explicit_run_id() {
+        let error = resolve_launch_run_id(2, None).unwrap_err();
+        assert!(error.to_string().contains("shared RUSTRAIN_RUN_ID"));
+    }
+
+    #[test]
+    fn multi_node_launch_requires_explicit_attempt_id() {
+        let error = resolve_launch_attempt_id(2, None).unwrap_err();
+        assert!(error.to_string().contains("shared RUSTRAIN_ATTEMPT_ID"));
+        assert_eq!(
+            resolve_launch_attempt_id(2, Some("attempt-7".into())).unwrap(),
+            "attempt-7"
+        );
+    }
+
+    #[test]
+    fn launch_ids_reject_path_components() {
+        for value in [".", "..", "../attempt"] {
+            let run_error = resolve_launch_run_id(1, Some(value.into())).unwrap_err();
+            assert!(run_error.to_string().contains("RUSTRAIN_RUN_ID"));
+            let attempt_error = resolve_launch_attempt_id(1, Some(value.into())).unwrap_err();
+            assert!(attempt_error.to_string().contains("RUSTRAIN_ATTEMPT_ID"));
+        }
+    }
+
+    #[test]
+    fn multi_node_rendezvous_observes_every_node_on_shared_storage() {
+        let temp = tempdir().unwrap();
+        let root = temp.path();
+        std::thread::scope(|scope| {
+            let handles = (0..2)
+                .map(|node_rank| {
+                    scope.spawn(move || {
+                        rendezvous_launch_nodes(
+                            root,
+                            2,
+                            4,
+                            node_rank,
+                            "127.0.0.1",
+                            29500,
+                            "shared-run",
+                            "attempt-7",
+                            Duration::from_secs(1),
+                        )
+                    })
+                })
+                .collect::<Vec<_>>();
+            for handle in handles {
+                handle.join().unwrap().unwrap();
+            }
+        });
+    }
+
+    #[test]
+    fn multi_node_rendezvous_reports_non_shared_storage() {
+        let temp = tempdir().unwrap();
+        let error = rendezvous_launch_nodes(
+            temp.path(),
+            2,
+            4,
+            0,
+            "127.0.0.1",
+            29500,
+            "shared-run",
+            "attempt-7",
+            Duration::from_millis(20),
+        )
+        .unwrap_err();
+        assert!(error.to_string().contains("shared storage"));
+    }
+
     #[test]
     fn launch_parses_visible_cuda_devices() {
         let devices = parse_visible_cuda_devices(Some("0, 2,GPU-abc".to_string()))
diff --git a/crates/rustrain-parallel/src/lib.rs b/crates/rustrain-parallel/src/lib.rs
index ad211184..e21308c5 100644
--- a/crates/rustrain-parallel/src/lib.rs
+++ b/crates/rustrain-parallel/src/lib.rs
@@ -2,3 +2,4 @@ pub mod dp_rank;
 pub mod launcher;
 pub mod parallel;
 pub mod parallel_modules;
+pub mod topology;
diff --git a/crates/rustrain-parallel/src/topology.rs b/crates/rustrain-parallel/src/topology.rs
new file mode 100644
index 00000000..0d8de3e2
--- /dev/null
+++ b/crates/rustrain-parallel/src/topology.rs
@@ -0,0 +1,561 @@
+//! Rank topology for Megatron-style orthogonal parallel groups.
+//!
+//! This module only describes process topology. It deliberately does not
+//! create NCCL process groups or shard model weights; those operations belong
+//! to the model/runtime that owns each collective. Keeping the rank mapping in
+//! one place prevents TP/PP/DP/EP/CP launchers from silently disagreeing.
+
+use std::{convert::TryInto, env, ops::Range};
+
+use anyhow::{Context, Result, anyhow, bail};
+use serde::{Deserialize, Serialize};
+
+use rustrain_core::runtime::ParallelConfig;
+
+/// Megatron's current default rank order for decoder groups.
+pub const DEFAULT_RANK_ORDER: &str = "tp-cp-ep-dp-pp";
+
+/// One orthogonal parallelism axis.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum ParallelAxis {
+    Tensor,
+    Pipeline,
+    Data,
+    Expert,
+    Context,
+}
+
+impl ParallelAxis {
+    pub const ALL: [Self; 5] = [
+        Self::Tensor,
+        Self::Context,
+        Self::Expert,
+        Self::Data,
+        Self::Pipeline,
+    ];
+
+    pub const fn name(self) -> &'static str {
+        match self {
+            Self::Tensor => "tp",
+            Self::Pipeline => "pp",
+            Self::Data => "dp",
+            Self::Expert => "ep",
+            Self::Context => "cp",
+        }
+    }
+
+    fn parse(raw: &str) -> Result<Self> {
+        match raw.trim().to_ascii_lowercase().as_str() {
+            "tp" | "tensor" => Ok(Self::Tensor),
+            "pp" | "pipeline" => Ok(Self::Pipeline),
+            "dp" | "data" => Ok(Self::Data),
+            "ep" | "expert" => Ok(Self::Expert),
+            "cp" | "context" => Ok(Self::Context),
+            other => bail!("unknown parallel axis '{other}' (expected tp/pp/dp/ep/cp)"),
+        }
+    }
+}
+
+/// Coordinates of one global rank in the five-dimensional topology.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+pub struct RankCoordinates {
+    pub tensor: usize,
+    pub pipeline: usize,
+    pub data: usize,
+    pub expert: usize,
+    pub context: usize,
+}
+
+impl RankCoordinates {
+    pub const ZERO: Self = Self {
+        tensor: 0,
+        pipeline: 0,
+        data: 0,
+        expert: 0,
+        context: 0,
+    };
+
+    pub const fn get(self, axis: ParallelAxis) -> usize {
+        match axis {
+            ParallelAxis::Tensor => self.tensor,
+            ParallelAxis::Pipeline => self.pipeline,
+            ParallelAxis::Data => self.data,
+            ParallelAxis::Expert => self.expert,
+            ParallelAxis::Context => self.context,
+        }
+    }
+
+    pub fn set(&mut self, axis: ParallelAxis, value: usize) {
+        match axis {
+            ParallelAxis::Tensor => self.tensor = value,
+            ParallelAxis::Pipeline => self.pipeline = value,
+            ParallelAxis::Data => self.data = value,
+            ParallelAxis::Expert => self.expert = value,
+            ParallelAxis::Context => self.context = value,
+        }
+    }
+}
+
+/// A validated orthogonal TP/PP/DP/EP/CP rank topology.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ParallelTopology {
+    sizes: RankCoordinates,
+    /// The first axis is the least-significant rank digit, matching
+    /// Megatron's `generate_masked_orthogonal_rank_groups` convention.
+    order: [ParallelAxis; 5],
+}
+
+impl ParallelTopology {
+    /// Construct a topology using the Megatron default order.
+    pub fn new(
+        tensor: usize,
+        pipeline: usize,
+        data: usize,
+        expert: usize,
+        context: usize,
+    ) -> Result<Self> {
+        Self::with_order(tensor, pipeline, data, expert, context, DEFAULT_RANK_ORDER)
+    }
+
+    /// Construct a topology with an explicit least-significant-first order.
+    pub fn with_order(
+        tensor: usize,
+        pipeline: usize,
+        data: usize,
+        expert: usize,
+        context: usize,
+        order: &str,
+    ) -> Result<Self> {
+        let sizes = RankCoordinates {
+            tensor,
+            pipeline,
+            data,
+            expert,
+            context,
+        };
+        for axis in ParallelAxis::ALL {
+            if sizes.get(axis) == 0 {
+                bail!("{} parallel size must be greater than zero", axis.name());
+            }
+        }
+
+        let parsed: Vec<ParallelAxis> = order
+            .split('-')
+            .filter(|token| !token.trim().is_empty())
+            .map(ParallelAxis::parse)
+            .collect::<Result<_>>()?;
+        if parsed.is_empty() || parsed.len() > ParallelAxis::ALL.len() {
+            bail!("parallel rank order must list each axis at most once (got '{order}')");
+        }
+        for axis in ParallelAxis::ALL {
+            let count = parsed
+                .iter()
+                .filter(|candidate| **candidate == axis)
+                .count();
+            if count > 1 {
+                bail!(
+                    "parallel rank order must contain {} at most once",
+                    axis.name()
+                );
+            }
+            if count == 0 && sizes.get(axis) != 1 {
+                bail!(
+                    "parallel rank order omits non-singleton {} axis (size {})",
+                    axis.name(),
+                    sizes.get(axis)
+                );
+            }
+        }
+        let mut order = parsed;
+        for axis in ParallelAxis::ALL {
+            if !order.contains(&axis) {
+                order.push(axis);
+            }
+        }
+        let order: [ParallelAxis; 5] = order
+            .try_into()
+            .map_err(|_| anyhow!("parallel rank order must contain five axes"))?;
+        Ok(Self { sizes, order })
+    }
+
+    /// Construct from the public runtime configuration and expected world size.
+    pub fn from_config(config: &ParallelConfig, world_size: usize) -> Result<Self> {
+        let topology = Self::new(
+            config.tensor_model_parallel_size,
+            config.pipeline_model_parallel_size,
+            config.data_parallel_size,
+            config.expert_model_parallel_size,
+            config.context_parallel_size,
+        )?;
+        topology.validate_world_size(world_size)?;
+        Ok(topology)
+    }
+
+    /// Read axis sizes from launcher environment variables.
+    ///
+    /// `TP_SIZE`, `PP_SIZE`, `DP_SIZE`, `EP_SIZE`, and `CP_SIZE` are accepted;
+    /// each also has a `RUSTRAIN_`-prefixed alias. When no axis is specified,
+    /// all ranks are treated as data-parallel replicas. If DP is omitted while
+    /// other axes are specified, it is inferred from `WORLD_SIZE`.
+    pub fn from_env() -> Result<Self> {
+        let world_size = parse_env_usize("WORLD_SIZE")?;
+        Self::from_env_with_world_size(world_size)
+    }
+
+    /// Read axis sizes from the launcher environment while taking world size
+    /// from the caller. This is used by the launcher before it spawns ranks,
+    /// when `WORLD_SIZE` is not yet present in the parent process.
+    pub fn from_env_with_world_size(world_size: usize) -> Result<Self> {
+        if world_size == 0 {
+            bail!("WORLD_SIZE must be greater than zero");
+        }
+        let mut values = [None; 5];
+        for (index, names) in [
+            ["TP_SIZE", "RUSTRAIN_TP_SIZE", "TENSOR_MODEL_PARALLEL_SIZE"],
+            [
+                "PP_SIZE",
+                "RUSTRAIN_PP_SIZE",
+                "PIPELINE_MODEL_PARALLEL_SIZE",
+            ],
+            ["DP_SIZE", "RUSTRAIN_DP_SIZE", "DATA_PARALLEL_SIZE"],
+            ["EP_SIZE", "RUSTRAIN_EP_SIZE", "EXPERT_MODEL_PARALLEL_SIZE"],
+            ["CP_SIZE", "RUSTRAIN_CP_SIZE", "CONTEXT_PARALLEL_SIZE"],
+        ]
+        .into_iter()
+        .enumerate()
+        {
+            values[index] = first_env_usize(&names)?;
+        }
+
+        let any_explicit = values.iter().any(Option::is_some);
+        let tensor = values[0].unwrap_or(1);
+        let pipeline = values[1].unwrap_or(1);
+        let expert = values[3].unwrap_or(1);
+        let context = values[4].unwrap_or(1);
+        let data = match values[2] {
+            Some(value) => value,
+            None if !any_explicit => world_size,
+            None => {
+                let model_parallel = tensor
+                    .checked_mul(pipeline)
+                    .and_then(|value| value.checked_mul(expert))
+                    .and_then(|value| value.checked_mul(context))
+                    .ok_or_else(|| anyhow!("parallel axis product overflowed usize"))?;
+                if model_parallel == 0 || world_size % model_parallel != 0 {
+                    bail!(
+                        "WORLD_SIZE={world_size} is not divisible by specified model-parallel product {model_parallel}; set DP_SIZE explicitly"
+                    );
+                }
+                world_size / model_parallel
+            }
+        };
+        let order = env::var("RUSTRAIN_PARALLEL_ORDER")
+            .or_else(|_| env::var("PARALLEL_ORDER"))
+            .unwrap_or_else(|_| DEFAULT_RANK_ORDER.to_string());
+        let topology = Self::with_order(tensor, pipeline, data, expert, context, &order)?;
+        topology.validate_world_size(world_size)?;
+        Ok(topology)
+    }
+
+    pub const fn order(&self) -> [ParallelAxis; 5] {
+        self.order
+    }
+
+    pub const fn sizes(&self) -> RankCoordinates {
+        self.sizes
+    }
+
+    pub const fn tensor_model_parallel_size(&self) -> usize {
+        self.sizes.tensor
+    }
+
+    pub const fn pipeline_model_parallel_size(&self) -> usize {
+        self.sizes.pipeline
+    }
+
+    pub const fn data_parallel_size(&self) -> usize {
+        self.sizes.data
+    }
+
+    pub const fn expert_model_parallel_size(&self) -> usize {
+        self.sizes.expert
+    }
+
+    pub const fn context_parallel_size(&self) -> usize {
+        self.sizes.context
+    }
+
+    pub fn world_size(&self) -> usize {
+        self.order
+            .into_iter()
+            .map(|axis| self.sizes.get(axis))
+            .product()
+    }
+
+    pub fn validate_world_size(&self, world_size: usize) -> Result<()> {
+        let expected = self.world_size();
+        if world_size != expected {
+            bail!(
+                "parallel topology expects world_size={expected} (tp={} pp={} dp={} ep={} cp={}), got {world_size}",
+                self.sizes.tensor,
+                self.sizes.pipeline,
+                self.sizes.data,
+                self.sizes.expert,
+                self.sizes.context,
+            );
+        }
+        Ok(())
+    }
+
+    /// Convert a global rank to its five local coordinates.
+    pub fn coordinates(&self, rank: usize) -> Result<RankCoordinates> {
+        if rank >= self.world_size() {
+            bail!(
+                "global rank {rank} is outside world_size {}",
+                self.world_size()
+            );
+        }
+        let mut remainder = rank;
+        let mut coordinates = RankCoordinates::ZERO;
+        for axis in self.order {
+            let size = self.sizes.get(axis);
+            coordinates.set(axis, remainder % size);
+            remainder /= size;
+        }
+        debug_assert_eq!(remainder, 0);
+        Ok(coordinates)
+    }
+
+    /// Convert local coordinates to a global rank.
+    pub fn rank(&self, coordinates: RankCoordinates) -> Result<usize> {
+        for axis in ParallelAxis::ALL {
+            if coordinates.get(axis) >= self.sizes.get(axis) {
+                bail!(
+                    "{} coordinate {} is outside size {}",
+                    axis.name(),
+                    coordinates.get(axis),
+                    self.sizes.get(axis)
+                );
+            }
+        }
+        let mut rank = 0usize;
+        let mut stride = 1usize;
+        for axis in self.order {
+            rank = rank
+                .checked_add(
+                    coordinates
+                        .get(axis)
+                        .checked_mul(stride)
+                        .ok_or_else(|| anyhow!("parallel rank calculation overflowed usize"))?,
+                )
+                .ok_or_else(|| anyhow!("parallel rank calculation overflowed usize"))?;
+            stride = stride
+                .checked_mul(self.sizes.get(axis))
+                .ok_or_else(|| anyhow!("parallel rank calculation overflowed usize"))?;
+        }
+        Ok(rank)
+    }
+
+    /// Return the global ranks in the process group for one axis.
+    pub fn group(&self, rank: usize, axis: ParallelAxis) -> Result<Vec<usize>> {
+        let coordinates = self.coordinates(rank)?;
+        let mut ranks = Vec::with_capacity(self.sizes.get(axis));
+        for local_rank in 0..self.sizes.get(axis) {
+            let mut member = coordinates;
+            member.set(axis, local_rank);
+            ranks.push(self.rank(member)?);
+        }
+        Ok(ranks)
+    }
+
+    pub fn tensor_group(&self, rank: usize) -> Result<Vec<usize>> {
+        self.group(rank, ParallelAxis::Tensor)
+    }
+
+    pub fn pipeline_group(&self, rank: usize) -> Result<Vec<usize>> {
+        self.group(rank, ParallelAxis::Pipeline)
+    }
+
+    pub fn data_group(&self, rank: usize) -> Result<Vec<usize>> {
+        self.group(rank, ParallelAxis::Data)
+    }
+
+    pub fn expert_group(&self, rank: usize) -> Result<Vec<usize>> {
+        self.group(rank, ParallelAxis::Expert)
+    }
+
+    pub fn context_group(&self, rank: usize) -> Result<Vec<usize>> {
+        self.group(rank, ParallelAxis::Context)
+    }
+
+    pub fn tensor_rank(&self, rank: usize) -> Result<usize> {
+        Ok(self.coordinates(rank)?.tensor)
+    }
+
+    pub fn pipeline_rank(&self, rank: usize) -> Result<usize> {
+        Ok(self.coordinates(rank)?.pipeline)
+    }
+
+    pub fn data_rank(&self, rank: usize) -> Result<usize> {
+        Ok(self.coordinates(rank)?.data)
+    }
+
+    pub fn expert_rank(&self, rank: usize) -> Result<usize> {
+        Ok(self.coordinates(rank)?.expert)
+    }
+
+    pub fn context_rank(&self, rank: usize) -> Result<usize> {
+        Ok(self.coordinates(rank)?.context)
+    }
+
+    pub fn is_first_pipeline_stage(&self, rank: usize) -> Result<bool> {
+        Ok(self.pipeline_rank(rank)? == 0)
+    }
+
+    pub fn is_last_pipeline_stage(&self, rank: usize) -> Result<bool> {
+        Ok(self.pipeline_rank(rank)? + 1 == self.pipeline_model_parallel_size())
+    }
+
+    /// Return the contiguous layer range assigned to a pipeline stage.
+    pub fn layer_range(&self, rank: usize, num_layers: usize) -> Result<Range<usize>> {
+        let stage = self.pipeline_rank(rank)?;
+        let stages = self.pipeline_model_parallel_size();
+        if num_layers < stages {
+            bail!("num_layers={num_layers} must be >= pipeline size {stages}");
+        }
+        let start = num_layers * stage / stages;
+        let end = num_layers * (stage + 1) / stages;
+        Ok(start..end)
+    }
+}
+
+fn parse_env_usize(name: &str) -> Result<usize> {
+    env::var(name)
+        .with_context(|| format!("{name} is not set"))?
+        .parse::<usize>()
+        .with_context(|| format!("{name} must be a positive integer"))
+}
+
+fn first_env_usize(names: &[&str]) -> Result<Option<usize>> {
+    let Some((name, raw)) = names
+        .iter()
+        .find_map(|name| env::var(name).ok().map(|raw| (*name, raw)))
+    else {
+        return Ok(None);
+    };
+    let value = raw
+        .parse::<usize>()
+        .with_context(|| format!("{name} must be a positive integer"))?;
+    if value == 0 {
+        bail!("{name} must be greater than zero");
+    }
+    Ok(Some(value))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn megatron_like_tp_dp_pp_order_round_trips_and_builds_groups() {
+        let topology = ParallelTopology::with_order(2, 4, 3, 1, 1, "tp-dp-pp").unwrap();
+        assert_eq!(topology.world_size(), 24);
+        let rank = topology
+            .rank(RankCoordinates {
+                tensor: 1,
+                pipeline: 2,
+                data: 1,
+                expert: 0,
+                context: 0,
+            })
+            .unwrap();
+        // tp is the least-significant digit: 1 + 1*2 + 2*2*3 = 15.
+        assert_eq!(rank, 15);
+        assert_eq!(topology.coordinates(rank).unwrap().tensor, 1);
+        assert_eq!(topology.coordinates(rank).unwrap().data, 1);
+        assert_eq!(topology.coordinates(rank).unwrap().pipeline, 2);
+        assert_eq!(topology.tensor_group(rank).unwrap(), vec![14, 15]);
+        assert_eq!(topology.data_group(rank).unwrap(), vec![13, 15, 17]);
+        assert_eq!(topology.pipeline_group(rank).unwrap(), vec![3, 9, 15, 21]);
+    }
+
+    #[test]
+    fn tp2_dp2_builds_orthogonal_groups_for_every_rank() {
+        let topology = ParallelTopology::new(2, 1, 2, 1, 1).unwrap();
+        let expected = [
+            (0, 0, vec![0, 1], vec![0, 2]),
+            (1, 0, vec![0, 1], vec![1, 3]),
+            (0, 1, vec![2, 3], vec![0, 2]),
+            (1, 1, vec![2, 3], vec![1, 3]),
+        ];
+        for (rank, (tp_rank, dp_rank, tp_group, dp_group)) in expected.into_iter().enumerate() {
+            assert_eq!(topology.tensor_rank(rank).unwrap(), tp_rank);
+            assert_eq!(topology.data_rank(rank).unwrap(), dp_rank);
+            assert_eq!(topology.tensor_group(rank).unwrap(), tp_group);
+            assert_eq!(topology.data_group(rank).unwrap(), dp_group);
+        }
+    }
+
+    #[test]
+    fn five_dimensional_default_order_matches_orthogonal_mapping() {
+        let topology = ParallelTopology::new(2, 2, 2, 2, 2).unwrap();
+        let coordinates = RankCoordinates {
+            tensor: 1,
+            pipeline: 0,
+            data: 1,
+            expert: 1,
+            context: 0,
+        };
+        // tp-cp-ep-dp-pp: 1 + 1*(2*2) + 1*(2*2*2) = 13.
+        assert_eq!(topology.rank(coordinates).unwrap(), 13);
+        assert_eq!(topology.coordinates(13).unwrap(), coordinates);
+        assert_eq!(topology.expert_group(13).unwrap(), vec![9, 13]);
+        assert_eq!(topology.context_group(13).unwrap(), vec![13, 15]);
+    }
+
+    #[test]
+    fn cp2_pp2_groups_match_native_communicator_colors() {
+        let topology = ParallelTopology::new(1, 2, 1, 1, 2).unwrap();
+        let expected = [
+            (0, 0, vec![0, 1], vec![0, 2]),
+            (1, 0, vec![0, 1], vec![1, 3]),
+            (0, 1, vec![2, 3], vec![0, 2]),
+            (1, 1, vec![2, 3], vec![1, 3]),
+        ];
+        for (rank, (cp_rank, pp_rank, cp_group, pp_group)) in
+            expected.into_iter().enumerate()
+        {
+            assert_eq!(topology.context_rank(rank).unwrap(), cp_rank);
+            assert_eq!(topology.pipeline_rank(rank).unwrap(), pp_rank);
+            assert_eq!(topology.context_group(rank).unwrap(), cp_group);
+            assert_eq!(topology.pipeline_group(rank).unwrap(), pp_group);
+            assert_eq!(
+                topology.context_group(rank).unwrap().into_iter().min(),
+                Some(if pp_rank == 0 { 0 } else { 2 })
+            );
+            assert_eq!(
+                topology.pipeline_group(rank).unwrap().into_iter().min(),
+                Some(cp_rank)
+            );
+        }
+    }
+
+    #[test]
+    fn rejects_invalid_order_and_world_size() {
+        assert!(ParallelTopology::with_order(1, 1, 1, 1, 1, "tp-unknown").is_err());
+        assert!(ParallelTopology::with_order(1, 1, 1, 1, 1, "tp-dp-dp-ep-cp").is_err());
+        let topology = ParallelTopology::new(2, 2, 1, 1, 1).unwrap();
+        assert!(topology.validate_world_size(3).is_err());
+        assert!(topology.coordinates(4).is_err());
+    }
+
+    #[test]
+    fn layer_ranges_are_contiguous_and_cover_the_model() {
+        let topology = ParallelTopology::new(1, 3, 1, 1, 1).unwrap();
+        assert_eq!(topology.layer_range(0, 10).unwrap(), 0..3);
+        assert_eq!(topology.layer_range(1, 10).unwrap(), 3..6);
+        assert_eq!(topology.layer_range(2, 10).unwrap(), 6..10);
+        assert!(topology.layer_range(0, 2).is_err());
+    }
+}
diff --git a/crates/rustrain-qwen3-6/build.rs b/crates/rustrain-qwen3-6/build.rs
index 243491cf..fa80d1bf 100644
--- a/crates/rustrain-qwen3-6/build.rs
+++ b/crates/rustrain-qwen3-6/build.rs
@@ -1,18 +1,6 @@
 // build.rs — Compile C++ Qwen3.6 kernels and link against libtorch.
-use std::path::PathBuf;
 use std::process::Command;
 
-fn which(cmd: &str) -> Option<String> {
-    std::process::Command::new("which")
-        .arg(cmd)
-        .output()
-        .ok()
-        .filter(|o| o.status.success())
-        .and_then(|o| String::from_utf8(o.stdout).ok())
-        .map(|s| s.trim().to_string())
-        .filter(|s| !s.is_empty())
-}
-
 /// Detect PyTorch's _GLIBCXX_USE_CXX11_ABI setting by running Python.
 /// Returns "1" or "0". Defaults to "1" if detection fails.
 fn detect_cxx11_abi() -> String {
@@ -23,7 +11,10 @@ fn detect_cxx11_abi() -> String {
     // Try python3 -c 'import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))'
     for py in &["python3", "python"] {
         if let Ok(out) = std::process::Command::new(py)
-            .args(["-c", "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))"])
+            .args([
+                "-c",
+                "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))",
+            ])
             .output()
         {
             if out.status.success() {
@@ -42,7 +33,15 @@ fn detect_cxx11_abi() -> String {
 fn main() {
     println!("cargo:rerun-if-env-changed=TORCH_INCLUDE_PATH");
     println!("cargo:rerun-if-env-changed=TORCH_LIB_PATH");
+    println!("cargo:rerun-if-env-changed=CUDA_HOME");
+    println!("cargo:rerun-if-env-changed=CUDA_INCLUDE_PATH");
+    println!("cargo:rerun-if-env-changed=CUDA_LIB_PATH");
+    println!("cargo:rerun-if-env-changed=NCCL_INCLUDE_PATH");
+    println!("cargo:rerun-if-env-changed=NCCL_LIB_PATH");
     println!("cargo:rerun-if-changed=kernels/qwen3_6_kernels.cpp");
+    println!("cargo:rerun-if-changed=kernels/delta_rule.cu");
+    println!("cargo:rerun-if-changed=kernels/delta_rule.cuh");
+    println!("cargo:rerun-if-changed=kernels/fused_kernels.cu");
     println!("cargo:rerun-if-changed=build.rs");
 
     let cxx11_abi = detect_cxx11_abi();
@@ -93,53 +92,152 @@ fn main() {
         }
     };
 
+    let torch_package_dir = std::path::Path::new(&torch_lib)
+        .parent()
+        .and_then(std::path::Path::parent);
+    let sibling_nccl = torch_package_dir.map(|dir| dir.join("nvidia/nccl"));
+    let nccl_include = std::env::var("NCCL_INCLUDE_PATH").unwrap_or_else(|_| {
+        let mut candidates = Vec::new();
+        if let Some(root) = &sibling_nccl {
+            candidates.push(root.join("include"));
+        }
+        candidates.extend([
+            std::path::PathBuf::from(
+                "/share/code/nolanho/mint-runtime-py31213/host-venv/lib/python3.12/site-packages/nvidia/nccl/include",
+            ),
+            std::path::PathBuf::from(
+                "/usr/local/lib/python3.13/dist-packages/nvidia/nccl/include",
+            ),
+            std::path::PathBuf::from("/usr/include"),
+        ]);
+        candidates
+            .into_iter()
+            .find(|path| path.join("nccl.h").exists())
+            .unwrap_or_else(|| std::path::PathBuf::from("/usr/include"))
+            .display()
+            .to_string()
+    });
+    let nccl_lib = std::env::var("NCCL_LIB_PATH").unwrap_or_else(|_| {
+        let mut candidates = Vec::new();
+        if let Some(root) = &sibling_nccl {
+            candidates.push(root.join("lib"));
+        }
+        candidates.extend([
+            std::path::PathBuf::from(
+                "/share/code/nolanho/mint-runtime-py31213/host-venv/lib/python3.12/site-packages/nvidia/nccl/lib",
+            ),
+            std::path::PathBuf::from(
+                "/usr/local/lib/python3.13/dist-packages/nvidia/nccl/lib",
+            ),
+            std::path::PathBuf::from("/usr/lib/x86_64-linux-gnu"),
+        ]);
+        candidates
+            .into_iter()
+            .find(|path| {
+                path.join("libnccl.so").exists() || path.join("libnccl.so.2").exists()
+            })
+            .unwrap_or_else(|| std::path::PathBuf::from("/usr/lib/x86_64-linux-gnu"))
+            .display()
+            .to_string()
+    });
+    let nccl_link = if std::path::Path::new(&nccl_lib).join("libnccl.so").exists() {
+        "-lnccl"
+    } else {
+        "-l:libnccl.so.2"
+    };
+
     let out_dir = std::env::var("OUT_DIR").unwrap_or_else(|_| "target/debug".to_string());
     let kernel_src = "kernels/qwen3_6_kernels.cpp";
     let output_lib = format!("{out_dir}/libqwen36_kernels.so");
+    // Never leave a previously built kernel at the current OUT_DIR after a
+    // failed rebuild. Runtime loading must not silently pick up stale code.
+    let _ = std::fs::remove_file(&output_lib);
 
     println!("cargo:warning=Compiling Qwen3.6 kernels: include={torch_include} lib={torch_lib}");
 
     let cuda_inc = std::env::var("CUDA_INCLUDE_PATH").unwrap_or_else(|_| {
         let candidates = [
             "/share/code/nolanho/pydeps/lora-research/nvidia/cu13/include",
+            "/usr/local/cuda-13/include",
             "/usr/local/cuda-13.0/include",
             "/usr/local/cuda/include",
         ];
         for c in &candidates {
-            if std::path::Path::new(&format!("{c}/cuda_runtime_api.h")).exists() {
+            if std::path::Path::new(&format!("{c}/cuda_runtime_api.h")).exists()
+                && std::path::Path::new(&format!("{c}/crt/host_defines.h")).exists()
+            {
                 return c.to_string();
             }
         }
         "/usr/local/cuda/include".to_string()
     });
+    let cuda_lib = std::env::var("CUDA_LIB_PATH").unwrap_or_else(|_| {
+        let mut candidates = Vec::new();
+        if let Ok(cuda_home) = std::env::var("CUDA_HOME") {
+            candidates.push(std::path::PathBuf::from(&cuda_home).join("lib64"));
+            candidates.push(std::path::PathBuf::from(cuda_home).join("targets/x86_64-linux/lib"));
+        }
+        if let Some(include_parent) = std::path::Path::new(&cuda_inc).parent() {
+            candidates.push(include_parent.join("lib64"));
+            candidates.push(include_parent.join("lib"));
+        }
+        candidates.extend([
+            std::path::PathBuf::from("/usr/local/cuda-13.0/lib64"),
+            std::path::PathBuf::from("/usr/local/cuda-13/lib64"),
+            std::path::PathBuf::from("/usr/local/cuda/lib64"),
+        ]);
+        candidates
+            .into_iter()
+            .find(|path| path.join("libcudart.so").exists())
+            .unwrap_or_else(|| std::path::PathBuf::from("/usr/local/cuda/lib64"))
+            .display()
+            .to_string()
+    });
 
-    let status = Command::new("g++")
+    let cpp_ok = Command::new("g++")
         .args([
-            "-shared", "-fPIC", "-std=c++17", "-O2",
-            &cxx11_flag,
-            "-fvisibility=default",
-            "-o", &output_lib, kernel_src,
-            &format!("-I{torch_include}"),
-            &format!("-I{torch_include}/ATen"),
-            &format!("-I{torch_include}/c10"),
-            &format!("-I{torch_include}/caffe2"),
-            &format!("-I{cuda_inc}"),
+            "-shared".to_string(),
+            "-fPIC".to_string(),
+            "-std=c++17".to_string(),
+            "-O2".to_string(),
+            cxx11_flag.clone(),
+            "-fvisibility=default".to_string(),
+            "-o".to_string(),
+            output_lib.clone(),
+            kernel_src.to_string(),
+            format!("-I{torch_include}"),
+            format!("-I{torch_include}/ATen"),
+            format!("-I{torch_include}/c10"),
+            format!("-I{torch_include}/caffe2"),
+            format!("-I{cuda_inc}"),
+            format!("-I{nccl_include}"),
+            format!("-L{torch_lib}"),
+            format!("-Wl,-rpath,{torch_lib}"),
+            format!("-L{nccl_lib}"),
+            format!("-Wl,-rpath,{nccl_lib}"),
+            "-Wl,--no-as-needed".to_string(),
+            "-Wl,--export-dynamic".to_string(),
+            "-ltorch".to_string(),
+            "-ltorch_cuda".to_string(),
+            "-ltorch_cpu".to_string(),
+            "-lc10".to_string(),
+            "-lc10_cuda".to_string(),
+            nccl_link.to_string(),
         ])
-        .args([
-            &format!("-L{torch_lib}"),
-            &format!("-Wl,-rpath,{torch_lib}"),
-            "-Wl,--no-as-needed",
-            "-Wl,--export-dynamic",
-            "-fvisibility=default",
-            "-ltorch", "-ltorch_cuda", "-ltorch_cpu", "-lc10", "-lc10_cuda",
-            "-lnccl",
-        ])
-        .status();
+        .status()
+        .map(|s| s.success())
+        .unwrap_or(false);
 
     // ── Compile CUDA kernels (.cu files) with nvcc ──
     let nvcc_path = {
+        let cuda_from_include = std::path::Path::new(&cuda_inc)
+            .parent()
+            .map(|home| home.join("bin/nvcc").display().to_string());
         let candidates = [
-            std::env::var("CUDA_HOME").ok().map(|h| format!("{h}/bin/nvcc")),
+            std::env::var("CUDA_HOME")
+                .ok()
+                .map(|h| format!("{h}/bin/nvcc")),
+            cuda_from_include,
             Some("nvcc".to_string()),
         ];
         let mut found = String::new();
@@ -149,7 +247,11 @@ fn main() {
                     found = c.clone();
                     break;
                 }
-                if std::process::Command::new(c).arg("--version").output().is_ok() {
+                if std::process::Command::new(c)
+                    .arg("--version")
+                    .output()
+                    .is_ok()
+                {
                     found = c.clone();
                     break;
                 }
@@ -158,7 +260,7 @@ fn main() {
         found
     };
 
-    if !nvcc_path.is_empty() {
+    let build_ok = if !nvcc_path.is_empty() {
         // CUDA files needing nvcc: delta_rule.cu (has __global__) + fused_kernels.cu (hand-written CUDA)
         let cu_files_nvcc = ["kernels/delta_rule.cu", "kernels/fused_kernels.cu"];
 
@@ -166,20 +268,29 @@ fn main() {
 
         // Compile CUDA kernels with nvcc
         for cu_file in &cu_files_nvcc {
-            if !std::path::Path::new(cu_file).exists() { continue; }
-            let obj_file = format!("{out_dir}/{}.o",
-                cu_file.replace("kernels/", "").replace(".cu", ""));
+            if !std::path::Path::new(cu_file).exists() {
+                continue;
+            }
+            let obj_file = format!(
+                "{out_dir}/{}.o",
+                cu_file.replace("kernels/", "").replace(".cu", "")
+            );
             let cu_status = Command::new(&nvcc_path)
                 .args([
-                    "-c", cu_file, "-o", &obj_file,
-                    "-O2", "-std=c++17",
+                    "-c",
+                    cu_file,
+                    "-o",
+                    &obj_file,
+                    "-O2",
+                    "-std=c++17",
                     &cxx11_flag,
                     &format!("-I{torch_include}"),
                     &format!("-I{torch_include}/ATen"),
                     &format!("-I{torch_include}/c10"),
                     &format!("-I{torch_include}/caffe2"),
                     &format!("-I{cuda_inc}"),
-                    "-Xcompiler", "-fPIC",
+                    "-Xcompiler",
+                    "-fPIC",
                 ])
                 .status();
             if cu_status.map(|s| s.success()).unwrap_or(false) {
@@ -187,46 +298,68 @@ fn main() {
             }
         }
 
-        // Re-link the .so with all .o files
-        if !obj_files.is_empty() {
+        // Re-link only when every required CUDA translation unit compiled.
+        if obj_files.len() == cu_files_nvcc.len() {
             let mut link_args = vec![
-                "-shared".to_string(), "-fPIC".to_string(), "-std=c++17".to_string(), "-O2".to_string(),
+                "-shared".to_string(),
+                "-fPIC".to_string(),
+                "-std=c++17".to_string(),
+                "-O2".to_string(),
                 cxx11_flag.clone(),
-                "-o".to_string(), output_lib.clone(), kernel_src.to_string(),
+                "-o".to_string(),
+                output_lib.clone(),
+                kernel_src.to_string(),
                 format!("-I{torch_include}"),
                 format!("-I{torch_include}/ATen"),
                 format!("-I{torch_include}/c10"),
                 format!("-I{torch_include}/caffe2"),
                 format!("-I{cuda_inc}"),
+                format!("-I{nccl_include}"),
                 format!("-L{torch_lib}"),
                 format!("-Wl,-rpath,{torch_lib}"),
-                format!("-L{cuda_inc}/../lib64"),
-                format!("-Wl,-rpath,{cuda_inc}/../lib64"),
+                format!("-L{nccl_lib}"),
+                format!("-Wl,-rpath,{nccl_lib}"),
+                format!("-L{cuda_lib}"),
+                format!("-Wl,-rpath,{cuda_lib}"),
                 "-Wl,--no-as-needed".to_string(),
-                "-ltorch".to_string(), "-ltorch_cuda".to_string(), "-ltorch_cpu".to_string(),
-                "-lc10".to_string(), "-lc10_cuda".to_string(), "-lcudart".to_string(), "-lnccl".to_string(),
+                "-ltorch".to_string(),
+                "-ltorch_cuda".to_string(),
+                "-ltorch_cpu".to_string(),
+                "-lc10".to_string(),
+                "-lc10_cuda".to_string(),
+                "-lcudart".to_string(),
+                nccl_link.to_string(),
             ];
             for obj in &obj_files {
                 link_args.push(obj.clone());
             }
-            let _ = Command::new("g++").args(&link_args).status();
+            cpp_ok
+                && Command::new("g++")
+                    .args(&link_args)
+                    .status()
+                    .map(|s| s.success())
+                    .unwrap_or(false)
+        } else {
+            false
         }
-    }
+    } else {
+        false
+    };
 
-    match status {
-        Ok(s) if s.success() => {
-            println!("cargo:rustc-link-search=native={torch_lib}");
-            println!("cargo:rustc-link-search=native={out_dir}");
-            println!("cargo:rustc-link-lib=dylib=qwen36_kernels");
-            println!("cargo:rustc-link-lib=dylib=c10");
-            println!("cargo:rustc-link-lib=dylib=torch");
-            println!("cargo:rustc-link-lib=dylib=torch_cpu");
-            println!("cargo:rustc-link-lib=dylib=torch_cuda");
-            println!("cargo:rustc-link-arg=-Wl,--no-as-needed");
-            println!("cargo:rustc-link-arg=-Wl,--allow-shlib-undefined");
-        }
-        _ => {
-            println!("cargo:warning=Failed to compile Qwen3.6 kernels, C++ path disabled");
-        }
+    if build_ok {
+        println!("cargo:rustc-link-search=native={torch_lib}");
+        println!("cargo:rustc-link-search=native={out_dir}");
+        println!("cargo:rustc-link-lib=dylib=qwen36_kernels");
+        println!("cargo:rustc-link-lib=dylib=c10");
+        println!("cargo:rustc-link-lib=dylib=torch");
+        println!("cargo:rustc-link-lib=dylib=torch_cpu");
+        println!("cargo:rustc-link-lib=dylib=torch_cuda");
+        println!("cargo:rustc-link-arg=-Wl,--no-as-needed");
+        println!("cargo:rustc-link-arg=-Wl,--allow-shlib-undefined");
+    } else {
+        let _ = std::fs::remove_file(&output_lib);
+        println!(
+            "cargo:warning=Failed to compile complete Qwen3.6 C++/CUDA kernels; native training path disabled"
+        );
     }
 }
diff --git a/crates/rustrain-qwen3-6/kernels/delta_rule.cu b/crates/rustrain-qwen3-6/kernels/delta_rule.cu
index 424d13e2..780f1215 100644
--- a/crates/rustrain-qwen3-6/kernels/delta_rule.cu
+++ b/crates/rustrain-qwen3-6/kernels/delta_rule.cu
@@ -4,29 +4,45 @@
 #include "delta_rule.cuh"
 
 // C-linkage wrapper for the forward host launcher
-extern "C" void cuda_gated_delta_rule(
+extern "C" int cuda_gated_delta_rule(
     const float* q, const float* k, const float* v,
     const float* g_exp, const float* beta,
     float* state, float* out, float* delta_buf,
-    int BH, int seq_len, int key_dim, int val_dim
+    float* state_checkpoints, int checkpoint_stride,
+    const int32_t* lengths, int heads_per_batch,
+    int BH, int seq_len, int key_dim, int val_dim, void* stream_ptr
 ) {
+    if (key_dim != DR_D_K || val_dim != DR_D_V || heads_per_batch <= 0) {
+        return -1;
+    }
     launch_gated_delta_rule(q, k, v, g_exp, beta, state, out, delta_buf,
-                            BH, seq_len, key_dim, val_dim);
+                            state_checkpoints, checkpoint_stride,
+                            lengths, heads_per_batch,
+                            BH, seq_len, key_dim, val_dim,
+                            reinterpret_cast<cudaStream_t>(stream_ptr));
+    return static_cast<int>(cudaGetLastError());
 }
 
 // C-linkage wrapper for the backward host launcher
-extern "C" void cuda_gated_delta_rule_backward(
+extern "C" int cuda_gated_delta_rule_backward(
     const float* q, const float* k, const float* v,
     const float* g_exp, const float* beta,
     const float* final_state,
     const float* delta_buf,
+    const float* state_checkpoints, int checkpoint_stride,
     const float* grad_out,
     float* grad_q, float* grad_k, float* grad_v,
     float* grad_g, float* grad_beta,
-    int BH, int seq_len, int key_dim, int val_dim
+    const int32_t* lengths, int heads_per_batch,
+    int BH, int seq_len, int key_dim, int val_dim, void* stream_ptr
 ) {
-    launch_gated_delta_rule_backward(q, k, v, g_exp, beta,
-        final_state, delta_buf, grad_out,
+    if (key_dim != DR_D_K || val_dim != DR_D_V || heads_per_batch <= 0) return -1;
+    int launch_status = launch_gated_delta_rule_backward(q, k, v, g_exp, beta,
+        final_state, delta_buf, state_checkpoints, checkpoint_stride, grad_out,
         grad_q, grad_k, grad_v, grad_g, grad_beta,
-        BH, seq_len, key_dim, val_dim);
+        lengths, heads_per_batch,
+        BH, seq_len, key_dim, val_dim,
+        reinterpret_cast<cudaStream_t>(stream_ptr));
+    if (launch_status != 0) return launch_status;
+    return static_cast<int>(cudaGetLastError());
 }
diff --git a/crates/rustrain-qwen3-6/kernels/delta_rule.cuh b/crates/rustrain-qwen3-6/kernels/delta_rule.cuh
index 712afdc0..0cab0b23 100644
--- a/crates/rustrain-qwen3-6/kernels/delta_rule.cuh
+++ b/crates/rustrain-qwen3-6/kernels/delta_rule.cuh
@@ -49,7 +49,11 @@
 
 #include <cuda_runtime.h>
 #include <cuda_fp16.h>
+#include <cstdint>
+#include <cstdlib>
 #include <cstdio>
+#include <cstring>
+#include <limits>
 
 // ──────────────────────────────────────────────────────────────────────
 // Constants
@@ -89,6 +93,10 @@ __global__ void gated_delta_rule_kernel(
     float* __restrict__ state,         // [BH, D_K, D_V]
     float* __restrict__ out,           // [BH, S, D_V]
     float* __restrict__ delta_buf,     // [BH, S, D_V] — saved for backward
+    float* __restrict__ state_checkpoints, // [BH, checkpoints, D_K, D_V]
+    int checkpoint_stride,             // <=0 disables boundary checkpoints
+    const int32_t* __restrict__ lengths, // [B], nullptr for dense sequences
+    int heads_per_batch,
     int S
 ) {
     const int bh = blockIdx.x;
@@ -111,6 +119,20 @@ __global__ void gated_delta_rule_kernel(
     const float* g_bh = g + bh * S;
     const float* beta_bh = beta + bh * S;
     float* out_bh = out + bh * S * DR_D_V;
+    const int valid_len = lengths == nullptr ? S :
+        max(0, min(S, lengths[bh / heads_per_batch]));
+
+    const int checkpoint_count = checkpoint_stride > 0
+        ? (S + checkpoint_stride - 1) / checkpoint_stride + 1
+        : 0;
+    if (state_checkpoints != nullptr) {
+        float* checkpoint = state_checkpoints +
+            static_cast<size_t>(bh) * checkpoint_count * DR_D_K * DR_D_V;
+        #pragma unroll
+        for (int i = tid; i < DR_D_K * DR_D_V; i += DR_THREADS)
+            checkpoint[i] = state_s[i];
+        __syncthreads();
+    }
 
     // Each thread processes column dv = tid
     // Access pattern: state_s[dk * DR_D_V + tid] for dk = 0..127
@@ -118,7 +140,7 @@ __global__ void gated_delta_rule_kernel(
     // Bank mapping: bank = (dk * DR_D_V + tid) % 32 = tid % 32 (since 128 % 32 == 0)
     // → All threads in a warp access different banks (tid 0..31 → banks 0..31) ✅ no conflict
 
-    for (int t = 0; t < S; t++) {
+    for (int t = 0; t < valid_len; t++) {
         const float g_t = g_bh[t];
         const float beta_t = beta_bh[t];
         const float* q_t = q_bh + t * DR_D_K;
@@ -131,7 +153,6 @@ __global__ void gated_delta_rule_kernel(
         for (int dk = 0; dk < DR_D_K; dk++) {
             state_s[dk * DR_D_V + tid] *= g_t;
         }
-        __syncthreads();
 
         // --- kv_mem = K[t] · S[:, dv] ---
         // Dot product of k_t[0..127] with state_s[0..127, tid]
@@ -144,8 +165,11 @@ __global__ void gated_delta_rule_kernel(
         // --- delta = (v - kv_mem) * beta ---
         float delta = (v_t - kv_mem) * beta_t;
 
-        // Save delta for backward pass
-        delta_buf[bh * S * DR_D_V + t * DR_D_V + tid] = delta;
+        // Save delta for the fused autograd backward. The reference backward
+        // path is selected explicitly and does not use this CUDA kernel.
+        if (delta_buf != nullptr) {
+            delta_buf[bh * S * DR_D_V + t * DR_D_V + tid] = delta;
+        }
 
         // --- State update: S[:, dv] += k_t * delta ---
         // Rank-1 update: each row gets k_t[dk] * delta added
@@ -153,7 +177,6 @@ __global__ void gated_delta_rule_kernel(
         for (int dk = 0; dk < DR_D_K; dk++) {
             state_s[dk * DR_D_V + tid] += k_t[dk] * delta;
         }
-        __syncthreads();
 
         // --- Output: out[t, dv] = Q[t] · S[:, dv] ---
         float out_val = 0.0f;
@@ -162,7 +185,37 @@ __global__ void gated_delta_rule_kernel(
             out_val += q_t[dk] * state_s[dk * DR_D_V + tid];
         }
         out_bh[t * DR_D_V + tid] = out_val;
-        __syncthreads();  // Ensure state visible before next token
+
+        if (state_checkpoints != nullptr &&
+            (t + 1) % checkpoint_stride == 0) {
+            const int checkpoint_index = (t + 1) / checkpoint_stride;
+            float* checkpoint = state_checkpoints +
+                (static_cast<size_t>(bh) * checkpoint_count + checkpoint_index) *
+                DR_D_K * DR_D_V;
+            #pragma unroll
+            for (int i = tid; i < DR_D_K * DR_D_V; i += DR_THREADS)
+                checkpoint[i] = state_s[i];
+        }
+    }
+
+    if (state_checkpoints != nullptr) {
+        const int completed = valid_len / checkpoint_stride;
+        const int first_unwritten = completed + 1;
+        const int num_chunks = checkpoint_count - 1;
+        for (int checkpoint_index = first_unwritten;
+            checkpoint_index <= num_chunks; ++checkpoint_index) {
+            float* checkpoint = state_checkpoints +
+                (static_cast<size_t>(bh) * checkpoint_count + checkpoint_index) *
+                DR_D_K * DR_D_V;
+            for (int i = tid; i < DR_D_K * DR_D_V; i += DR_THREADS)
+                checkpoint[i] = state_s[i];
+        }
+    }
+
+    for (int t = valid_len; t < S; ++t) {
+        out_bh[t * DR_D_V + tid] = 0.0f;
+        if (delta_buf != nullptr)
+            delta_buf[bh * S * DR_D_V + t * DR_D_V + tid] = 0.0f;
     }
 
     // Write back state to global memory
@@ -181,10 +234,12 @@ inline void launch_gated_delta_rule(
     const float* q, const float* k, const float* v,
     const float* g_exp, const float* beta,
     float* state, float* out, float* delta_buf,
+    float* state_checkpoints, int checkpoint_stride,
+    const int32_t* lengths, int heads_per_batch,
     int BH, int seq_len, int key_dim, int val_dim,
-    cudaStream_t stream = 0
+    cudaStream_t stream
 ) {
-    if (key_dim != DR_D_K || val_dim != DR_D_V) {
+    if (key_dim != DR_D_K || val_dim != DR_D_V || heads_per_batch <= 0) {
         fprintf(stderr, "[delta_rule] ERROR: D_K=%d or D_V=%d mismatch (expected %d/%d)\n",
                 key_dim, val_dim, DR_D_K, DR_D_V);
         return;
@@ -200,7 +255,9 @@ inline void launch_gated_delta_rule(
         cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
 
     gated_delta_rule_kernel<<<grid, block, smem_size, stream>>>(
-        q, k, v, g_exp, beta, state, out, delta_buf, seq_len
+        q, k, v, g_exp, beta, state, out, delta_buf,
+        state_checkpoints, checkpoint_stride,
+        lengths, heads_per_batch, seq_len
     );
 }
 
@@ -262,7 +319,6 @@ __global__ void gated_delta_rule_backward_kernel(
 
     const float* q_bh = q + bh * S * DR_D_K;
     const float* k_bh = k + bh * S * DR_D_K;
-    const float* v_bh = v + bh * S * DR_D_V;
     const float* g_bh = g_exp + bh * S;
     const float* beta_bh = beta + bh * S;
     const float* go_bh = grad_out + bh * S * DR_D_V;
@@ -279,7 +335,6 @@ __global__ void gated_delta_rule_backward_kernel(
         const float beta_t = beta_bh[t];
         const float* q_t = q_bh + t * DR_D_K;
         const float* k_t = k_bh + t * DR_D_K;
-        const float v_t = v_bh[t * DR_D_V + tid];
         const float go_t = go_bh[t * DR_D_V + tid];
         const float delta_t = delta_buf[bh * S * DR_D_V + t * DR_D_V + tid];
 
@@ -337,7 +392,7 @@ __global__ void gated_delta_rule_backward_kernel(
             // Actually beta is [BH, S], not per-dv. Need reduction.
             // For now, store per-dv and reduce later.
         }
-        gb_bh[t] = 0.0f;  // placeholder — needs warp reduction
+        gb_bh[t] = 0.0f;  // legacy kernel; not referenced by the host launcher
 
         // --- Step 5: grad_k from kv ---
         // d(kv)/d(k) = S_before_g → grad_k_kv = S_before_g · grad_kv
@@ -360,43 +415,870 @@ __global__ void gated_delta_rule_backward_kernel(
         // --- Step 7: grad_g = sum(grad_S_before_decay * S_before) ---
         // This needs S_before = state_s / g[t], and grad_S_before_decay
         // For simplicity, approximate grad_g as 0 (g is exp of a_log, small gradient)
-        gg_bh[t] = 0.0f;  // TODO: compute properly
+        gg_bh[t] = 0.0f;  // legacy kernel; not referenced by the host launcher
+
+        __syncthreads();
+    }
+}
+
+// Correct reverse-mode recurrence. The older kernel above is retained only as
+// a historical reference and is not used by any launcher; this kernel computes
+// all q/k/v/g/beta gradients.
+// The forward recurrence is:
+//   R = g * S_prev, kv = k^T R, delta = beta * (v - kv),
+//   S = R + k outer delta, out = q^T S.
+template <bool kFuseRecurrence>
+__global__ void gated_delta_rule_backward_kernel_correct(
+    const float* __restrict__ q,
+    const float* __restrict__ k,
+    const float* __restrict__ v,
+    const float* __restrict__ g_exp,
+    const float* __restrict__ beta,
+    const float* __restrict__ final_state,
+    const float* __restrict__ delta_buf,
+    const float* __restrict__ state_checkpoints,
+    int checkpoint_stride,
+    const float* __restrict__ grad_out,
+    float* __restrict__ grad_q,
+    float* __restrict__ grad_k,
+    float* __restrict__ grad_v,
+    float* __restrict__ grad_g,
+    float* __restrict__ grad_beta,
+    const int32_t* __restrict__ lengths,
+    int heads_per_batch,
+    int S
+) {
+    const int bh = blockIdx.x;
+    const int tid = threadIdx.x;
+    const int lane = tid & 31;
+    const int warp = tid >> 5;
+
+    extern __shared__ float smem[];
+    float* state_s = smem;  // [D_K, D_V] = S_t while entering each step
+    float* grad_s = state_s + DR_D_K * DR_D_V;
+    // Four warp partials for reductions over D_V, followed by two scalar
+    // reductions. This avoids atomics and keeps the 128-thread block intact.
+    float* reduce_q = grad_s + DR_D_K * DR_D_V;
+    float* reduce_k = reduce_q + 4 * DR_D_K;
+    float* reduce_beta = reduce_k + 4 * DR_D_K;
+    float* reduce_g = reduce_beta + 4;
+
+    const float* state_g = final_state + bh * DR_D_K * DR_D_V;
+    for (int i = tid; i < DR_D_K * DR_D_V; i += DR_THREADS) {
+        state_s[i] = state_g[i];
+        grad_s[i] = 0.0f;
+    }
+    __syncthreads();
+
+    const float* q_bh = q + bh * S * DR_D_K;
+    const float* k_bh = k + bh * S * DR_D_K;
+    const float* v_bh = v + bh * S * DR_D_V;
+    const float* g_bh = g_exp + bh * S;
+    const float* beta_bh = beta + bh * S;
+    const float* go_bh = grad_out + bh * S * DR_D_V;
+    float* gq_bh = grad_q + bh * S * DR_D_K;
+    float* gk_bh = grad_k + bh * S * DR_D_K;
+    float* gv_bh = grad_v + bh * S * DR_D_V;
+    float* gg_bh = grad_g + bh * S;
+    float* gb_bh = grad_beta + bh * S;
+    const int valid_len = lengths == nullptr ? S :
+        max(0, min(S, lengths[bh / heads_per_batch]));
+
+    for (int t = valid_len; t < S; ++t) {
+        if (tid < DR_D_K) {
+            gq_bh[t * DR_D_K + tid] = 0.0f;
+            gk_bh[t * DR_D_K + tid] = 0.0f;
+        }
+        gv_bh[t * DR_D_V + tid] = 0.0f;
+        if (tid == 0) {
+            gg_bh[t] = 0.0f;
+            gb_bh[t] = 0.0f;
+        }
+    }
+    __syncthreads();
+
+    const int checkpoint_count = checkpoint_stride > 0
+        ? (S + checkpoint_stride - 1) / checkpoint_stride + 1
+        : 0;
+    const int num_chunks = state_checkpoints != nullptr
+        ? checkpoint_count - 1
+        : 1;
+    for (int chunk = num_chunks - 1; chunk >= 0; --chunk) {
+        const int chunk_start = state_checkpoints != nullptr
+            ? chunk * checkpoint_stride : 0;
+        const int chunk_end = state_checkpoints != nullptr
+            ? min(S, (chunk + 1) * checkpoint_stride) : S;
+        if (state_checkpoints != nullptr) {
+            const float* checkpoint = state_checkpoints +
+                (static_cast<size_t>(bh) * checkpoint_count + chunk + 1) *
+                DR_D_K * DR_D_V;
+            for (int i = tid; i < DR_D_K * DR_D_V; i += DR_THREADS)
+                state_s[i] = checkpoint[i];
+            __syncthreads();
+        }
+        const int reverse_end = min(valid_len, chunk_end);
+        for (int t = reverse_end - 1; t >= chunk_start; --t) {
+            const float* q_t = q_bh + t * DR_D_K;
+            const float* k_t = k_bh + t * DR_D_K;
+            const float g_t = g_bh[t];
+            const float beta_t = beta_bh[t];
+            const float go_t = go_bh[t * DR_D_V + tid];
+            const float delta_t =
+                delta_buf[bh * S * DR_D_V + t * DR_D_V + tid];
+
+            // Add the direct output contribution to dS_t and reduce dQ_t over
+            // value columns while state_s still contains the post-update state.
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                const float s_after = state_s[idx];
+                grad_s[idx] += q_t[dk] * go_t;
+                if constexpr (kFuseRecurrence)
+                    state_s[idx] = s_after - k_t[dk] * delta_t;
+                float q_part = s_after * go_t;
+                for (int off = 16; off > 0; off >>= 1)
+                    q_part += __shfl_down_sync(0xffffffff, q_part, off);
+                if (lane == 0) reduce_q[warp * DR_D_K + dk] = q_part;
+            }
+            __syncthreads();
+            if (tid < DR_D_K) {
+                gq_bh[t * DR_D_K + tid] =
+                    reduce_q[tid] + reduce_q[DR_D_K + tid] +
+                    reduce_q[2 * DR_D_K + tid] +
+                    reduce_q[3 * DR_D_K + tid];
+            }
+            if constexpr (!kFuseRecurrence) __syncthreads();
+
+            // Undo S_t = R_t + k outer delta, leaving R_t = g_t*S_prev.
+            if constexpr (!kFuseRecurrence) {
+                for (int dk = 0; dk < DR_D_K; ++dk)
+                    state_s[dk * DR_D_V + tid] -= k_t[dk] * delta_t;
+                __syncthreads();
+            }
+
+            // ddelta = dS_t^T k and h = ddelta * beta. The same value-column
+            // thread computes dV and one partial for dBeta/dG.
+            float gdelta = 0.0f;
+            float kv_mem = 0.0f;
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                gdelta += grad_s[idx] * k_t[dk];
+                kv_mem += state_s[idx] * k_t[dk];
+            }
+            const float h = gdelta * beta_t;
+            gv_bh[t * DR_D_V + tid] = h;
+            const float beta_part =
+                gdelta * (v_bh[t * DR_D_V + tid] - kv_mem);
+            const float safe_g = fmaxf(g_t, 1.0e-8f);
+            float g_part = 0.0f;
+
+            // dR = dS - k outer h. Reduce dK over value columns and update
+            // dS_prev.
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                const float r = state_s[idx];
+                const float s_prev = r / safe_g;
+                const float grad_r = grad_s[idx] - k_t[dk] * h;
+                const float k_part = grad_s[idx] * delta_t - r * h;
+                g_part += grad_r * s_prev;
+                float reduced_k = k_part;
+                for (int off = 16; off > 0; off >>= 1)
+                    reduced_k += __shfl_down_sync(
+                        0xffffffff, reduced_k, off);
+                if (lane == 0)
+                    reduce_k[warp * DR_D_K + dk] = reduced_k;
+                grad_s[idx] = grad_r * g_t;
+                // The next reverse iteration starts from S_prev, not R_t.
+                state_s[idx] = s_prev;
+            }
+
+            float reduced_beta = beta_part;
+            float reduced_g = g_part;
+            for (int off = 16; off > 0; off >>= 1) {
+                reduced_beta += __shfl_down_sync(
+                    0xffffffff, reduced_beta, off);
+                reduced_g += __shfl_down_sync(0xffffffff, reduced_g, off);
+            }
+            if (lane == 0) {
+                reduce_beta[warp] = reduced_beta;
+                reduce_g[warp] = reduced_g;
+            }
+            __syncthreads();
+            if (tid < DR_D_K) {
+                gk_bh[t * DR_D_K + tid] =
+                    reduce_k[tid] + reduce_k[DR_D_K + tid] +
+                    reduce_k[2 * DR_D_K + tid] +
+                    reduce_k[3 * DR_D_K + tid];
+            }
+            if (tid == 0) {
+                gb_bh[t] = reduce_beta[0] + reduce_beta[1] +
+                           reduce_beta[2] + reduce_beta[3];
+                gg_bh[t] = reduce_g[0] + reduce_g[1] +
+                           reduce_g[2] + reduce_g[3];
+            }
+            if constexpr (!kFuseRecurrence) __syncthreads();
+        }
+    }
+}
+
+// Chunkwise backward splits the recurrent dependency from the expensive
+// parameter-gradient work. The boundary pass computes only dS at exact state
+// checkpoint boundaries. Independent chunk CTAs then replay the full backward
+// from their state/dS end boundaries, increasing parallelism when BH alone does
+// not fill the GPU.
+__global__ void gated_delta_rule_backward_boundaries_kernel(
+    const float* __restrict__ q,
+    const float* __restrict__ k,
+    const float* __restrict__ g_exp,
+    const float* __restrict__ beta,
+    const float* __restrict__ grad_out,
+    float* __restrict__ grad_boundaries,
+    const int32_t* __restrict__ lengths,
+    int heads_per_batch,
+    int S,
+    int checkpoint_stride,
+    int checkpoint_count
+) {
+    const int bh = blockIdx.x;
+    const int tid = threadIdx.x;
+    constexpr int state_elements = DR_D_K * DR_D_V;
+
+    extern __shared__ float grad_s[];
+    for (int i = tid; i < state_elements; i += DR_THREADS)
+        grad_s[i] = 0.0f;
+
+    const float* q_bh = q + bh * S * DR_D_K;
+    const float* k_bh = k + bh * S * DR_D_K;
+    const float* g_bh = g_exp + bh * S;
+    const float* beta_bh = beta + bh * S;
+    const float* go_bh = grad_out + bh * S * DR_D_V;
+    const int valid_len = lengths == nullptr ? S :
+        max(0, min(S, lengths[bh / heads_per_batch]));
+    const int num_chunks = checkpoint_count - 1;
+
+    float* final_boundary = grad_boundaries +
+        (static_cast<size_t>(bh) * checkpoint_count + num_chunks) *
+            state_elements;
+    for (int i = tid; i < state_elements; i += DR_THREADS)
+        final_boundary[i] = 0.0f;
+
+    // Each thread exclusively owns one value column of dS, so this recurrence
+    // needs no inter-thread synchronization inside the token loop.
+    for (int chunk = num_chunks - 1; chunk >= 0; --chunk) {
+        const int chunk_start = chunk * checkpoint_stride;
+        const int reverse_end = min(valid_len, (chunk + 1) * checkpoint_stride);
+        for (int t = reverse_end - 1; t >= chunk_start; --t) {
+            const float* q_t = q_bh + t * DR_D_K;
+            const float* k_t = k_bh + t * DR_D_K;
+            const float go_t = go_bh[t * DR_D_V + tid];
+            float gdelta = 0.0f;
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                grad_s[idx] += q_t[dk] * go_t;
+                gdelta += grad_s[idx] * k_t[dk];
+            }
+            const float h = gdelta * beta_bh[t];
+            const float g_t = g_bh[t];
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                grad_s[idx] = (grad_s[idx] - k_t[dk] * h) * g_t;
+            }
+        }
+
+        float* boundary = grad_boundaries +
+            (static_cast<size_t>(bh) * checkpoint_count + chunk) *
+                state_elements;
+        for (int i = tid; i < state_elements; i += DR_THREADS)
+            boundary[i] = grad_s[i];
+    }
+}
+
+__global__ void gated_delta_rule_backward_chunks_kernel(
+    const float* __restrict__ q,
+    const float* __restrict__ k,
+    const float* __restrict__ v,
+    const float* __restrict__ g_exp,
+    const float* __restrict__ beta,
+    const float* __restrict__ delta_buf,
+    const float* __restrict__ state_checkpoints,
+    const float* __restrict__ grad_boundaries,
+    int checkpoint_stride,
+    int checkpoint_count,
+    const float* __restrict__ grad_out,
+    float* __restrict__ grad_q,
+    float* __restrict__ grad_k,
+    float* __restrict__ grad_v,
+    float* __restrict__ grad_g,
+    float* __restrict__ grad_beta,
+    const int32_t* __restrict__ lengths,
+    int heads_per_batch,
+    int S
+) {
+    const int num_chunks = checkpoint_count - 1;
+    const int bh = blockIdx.x / num_chunks;
+    const int chunk = blockIdx.x - bh * num_chunks;
+    const int tid = threadIdx.x;
+    const int lane = tid & 31;
+    const int warp = tid >> 5;
+    constexpr int state_elements = DR_D_K * DR_D_V;
+
+    extern __shared__ float smem[];
+    float* state_s = smem;
+    float* grad_s = state_s + state_elements;
+    float* reduce_q = grad_s + state_elements;
+    float* reduce_k = reduce_q + 4 * DR_D_K;
+    float* reduce_beta = reduce_k + 4 * DR_D_K;
+    float* reduce_g = reduce_beta + 4;
+
+    const float* state_boundary = state_checkpoints +
+        (static_cast<size_t>(bh) * checkpoint_count + chunk + 1) *
+            state_elements;
+    const float* grad_boundary = grad_boundaries +
+        (static_cast<size_t>(bh) * checkpoint_count + chunk + 1) *
+            state_elements;
+    for (int i = tid; i < state_elements; i += DR_THREADS) {
+        state_s[i] = state_boundary[i];
+        grad_s[i] = grad_boundary[i];
+    }
+    __syncthreads();
+
+    const float* q_bh = q + bh * S * DR_D_K;
+    const float* k_bh = k + bh * S * DR_D_K;
+    const float* v_bh = v + bh * S * DR_D_V;
+    const float* g_bh = g_exp + bh * S;
+    const float* beta_bh = beta + bh * S;
+    const float* go_bh = grad_out + bh * S * DR_D_V;
+    float* gq_bh = grad_q + bh * S * DR_D_K;
+    float* gk_bh = grad_k + bh * S * DR_D_K;
+    float* gv_bh = grad_v + bh * S * DR_D_V;
+    float* gg_bh = grad_g + bh * S;
+    float* gb_bh = grad_beta + bh * S;
+    const int valid_len = lengths == nullptr ? S :
+        max(0, min(S, lengths[bh / heads_per_batch]));
+    const int chunk_start = chunk * checkpoint_stride;
+    const int chunk_end = min(S, (chunk + 1) * checkpoint_stride);
+    const int reverse_end = min(valid_len, chunk_end);
 
+    for (int t = max(valid_len, chunk_start); t < chunk_end; ++t) {
+        if (tid < DR_D_K) {
+            gq_bh[t * DR_D_K + tid] = 0.0f;
+            gk_bh[t * DR_D_K + tid] = 0.0f;
+        }
+        gv_bh[t * DR_D_V + tid] = 0.0f;
+        if (tid == 0) {
+            gg_bh[t] = 0.0f;
+            gb_bh[t] = 0.0f;
+        }
+    }
+    __syncthreads();
+
+    for (int t = reverse_end - 1; t >= chunk_start; --t) {
+        const float* q_t = q_bh + t * DR_D_K;
+        const float* k_t = k_bh + t * DR_D_K;
+        const float g_t = g_bh[t];
+        const float beta_t = beta_bh[t];
+        const float go_t = go_bh[t * DR_D_V + tid];
+        const float delta_t =
+            delta_buf[bh * S * DR_D_V + t * DR_D_V + tid];
+
+        for (int dk = 0; dk < DR_D_K; ++dk) {
+            const int idx = dk * DR_D_V + tid;
+            const float s_after = state_s[idx];
+            grad_s[idx] += q_t[dk] * go_t;
+            state_s[idx] = s_after - k_t[dk] * delta_t;
+            float q_part = s_after * go_t;
+            for (int off = 16; off > 0; off >>= 1)
+                q_part += __shfl_down_sync(0xffffffff, q_part, off);
+            if (lane == 0) reduce_q[warp * DR_D_K + dk] = q_part;
+        }
+        __syncthreads();
+        if (tid < DR_D_K) {
+            gq_bh[t * DR_D_K + tid] =
+                reduce_q[tid] + reduce_q[DR_D_K + tid] +
+                reduce_q[2 * DR_D_K + tid] +
+                reduce_q[3 * DR_D_K + tid];
+        }
+
+        float gdelta = 0.0f;
+        float kv_mem = 0.0f;
+        for (int dk = 0; dk < DR_D_K; ++dk) {
+            const int idx = dk * DR_D_V + tid;
+            gdelta += grad_s[idx] * k_t[dk];
+            kv_mem += state_s[idx] * k_t[dk];
+        }
+        const float h = gdelta * beta_t;
+        gv_bh[t * DR_D_V + tid] = h;
+        const float beta_part =
+            gdelta * (v_bh[t * DR_D_V + tid] - kv_mem);
+        const float safe_g = fmaxf(g_t, 1.0e-8f);
+        float g_part = 0.0f;
+
+        for (int dk = 0; dk < DR_D_K; ++dk) {
+            const int idx = dk * DR_D_V + tid;
+            const float r = state_s[idx];
+            const float s_prev = r / safe_g;
+            const float grad_r = grad_s[idx] - k_t[dk] * h;
+            const float k_part = grad_s[idx] * delta_t - r * h;
+            g_part += grad_r * s_prev;
+            float reduced_k = k_part;
+            for (int off = 16; off > 0; off >>= 1)
+                reduced_k += __shfl_down_sync(
+                    0xffffffff, reduced_k, off);
+            if (lane == 0)
+                reduce_k[warp * DR_D_K + dk] = reduced_k;
+            grad_s[idx] = grad_r * g_t;
+            state_s[idx] = s_prev;
+        }
+
+        float reduced_beta = beta_part;
+        float reduced_g = g_part;
+        for (int off = 16; off > 0; off >>= 1) {
+            reduced_beta += __shfl_down_sync(
+                0xffffffff, reduced_beta, off);
+            reduced_g += __shfl_down_sync(0xffffffff, reduced_g, off);
+        }
+        if (lane == 0) {
+            reduce_beta[warp] = reduced_beta;
+            reduce_g[warp] = reduced_g;
+        }
+        __syncthreads();
+        if (tid < DR_D_K) {
+            gk_bh[t * DR_D_K + tid] =
+                reduce_k[tid] + reduce_k[DR_D_K + tid] +
+                reduce_k[2 * DR_D_K + tid] +
+                reduce_k[3 * DR_D_K + tid];
+        }
+        if (tid == 0) {
+            gb_bh[t] = reduce_beta[0] + reduce_beta[1] +
+                       reduce_beta[2] + reduce_beta[3];
+            gg_bh[t] = reduce_g[0] + reduce_g[1] +
+                       reduce_g[2] + reduce_g[3];
+        }
         __syncthreads();
     }
 }
 
+// Stable reverse-mode recurrence. Each CTA owns one batch-head. For every
+// reverse chunk it reloads the exact state at the chunk start, replays the
+// forward recurrence once into a compact per-CTA global workspace, then reads
+// exact S_{t-1}/S_t pairs while computing gradients. Unlike the legacy kernels,
+// this never reconstructs S_{t-1} by subtracting the rank-1 update and dividing
+// by g_t, which is ill-conditioned when the learned decay is tiny.
+__global__ void gated_delta_rule_backward_replay_kernel(
+    const float* __restrict__ q,
+    const float* __restrict__ k,
+    const float* __restrict__ v,
+    const float* __restrict__ g_exp,
+    const float* __restrict__ beta,
+    const float* __restrict__ delta_buf,
+    const float* __restrict__ state_checkpoints,
+    int checkpoint_stride,
+    float* __restrict__ replay_states,
+    const float* __restrict__ grad_out,
+    float* __restrict__ grad_q,
+    float* __restrict__ grad_k,
+    float* __restrict__ grad_v,
+    float* __restrict__ grad_g,
+    float* __restrict__ grad_beta,
+    const int32_t* __restrict__ lengths,
+    int heads_per_batch,
+    int S
+) {
+    const int bh = blockIdx.x;
+    const int tid = threadIdx.x;
+    const int lane = tid & 31;
+    const int warp = tid >> 5;
+    constexpr int state_elements = DR_D_K * DR_D_V;
+
+    extern __shared__ float smem[];
+    float* state_s = smem;
+    float* grad_s = state_s + state_elements;
+    float* reduce_q = grad_s + state_elements;
+    float* reduce_k = reduce_q + 4 * DR_D_K;
+    float* reduce_beta = reduce_k + 4 * DR_D_K;
+    float* reduce_g = reduce_beta + 4;
+
+    for (int i = tid; i < state_elements; i += DR_THREADS)
+        grad_s[i] = 0.0f;
+
+    const float* q_bh = q + bh * S * DR_D_K;
+    const float* k_bh = k + bh * S * DR_D_K;
+    const float* v_bh = v + bh * S * DR_D_V;
+    const float* g_bh = g_exp + bh * S;
+    const float* beta_bh = beta + bh * S;
+    const float* go_bh = grad_out + bh * S * DR_D_V;
+    float* gq_bh = grad_q + bh * S * DR_D_K;
+    float* gk_bh = grad_k + bh * S * DR_D_K;
+    float* gv_bh = grad_v + bh * S * DR_D_V;
+    float* gg_bh = grad_g + bh * S;
+    float* gb_bh = grad_beta + bh * S;
+    const int valid_len = lengths == nullptr ? S :
+        max(0, min(S, lengths[bh / heads_per_batch]));
+
+    for (int t = valid_len; t < S; ++t) {
+        if (tid < DR_D_K) {
+            gq_bh[t * DR_D_K + tid] = 0.0f;
+            gk_bh[t * DR_D_K + tid] = 0.0f;
+        }
+        gv_bh[t * DR_D_V + tid] = 0.0f;
+        if (tid == 0) {
+            gg_bh[t] = 0.0f;
+            gb_bh[t] = 0.0f;
+        }
+    }
+    __syncthreads();
+
+    const int checkpoint_count =
+        (S + checkpoint_stride - 1) / checkpoint_stride + 1;
+    const int num_chunks = checkpoint_count - 1;
+    float* replay_bh = replay_states +
+        static_cast<size_t>(bh) * (checkpoint_stride + 1) * state_elements;
+
+    for (int chunk = num_chunks - 1; chunk >= 0; --chunk) {
+        const int chunk_start = chunk * checkpoint_stride;
+        const int chunk_end = min(S, (chunk + 1) * checkpoint_stride);
+        const int reverse_end = min(valid_len, chunk_end);
+        if (reverse_end <= chunk_start) continue;
+
+        const float* chunk_start_state = state_checkpoints +
+            (static_cast<size_t>(bh) * checkpoint_count + chunk) *
+                state_elements;
+        for (int i = tid; i < state_elements; i += DR_THREADS) {
+            const float initial = chunk_start_state[i];
+            state_s[i] = initial;
+            replay_bh[i] = initial;
+        }
+        __syncthreads();
+
+        // Replay the exact forward update order and retain only this chunk's
+        // states. delta_buf is the forward-computed innovation, so replay does
+        // not introduce a second beta/KV evaluation or extra rounding there.
+        for (int t = chunk_start; t < reverse_end; ++t) {
+            const float g_t = g_bh[t];
+            const float* k_t = k_bh + t * DR_D_K;
+            const float delta_t =
+                delta_buf[bh * S * DR_D_V + t * DR_D_V + tid];
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                state_s[idx] *= g_t;
+            }
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                state_s[idx] += k_t[dk] * delta_t;
+            }
+            float* replay_slot = replay_bh +
+                static_cast<size_t>(t - chunk_start + 1) * state_elements;
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                replay_slot[idx] = state_s[idx];
+            }
+        }
+
+        for (int t = reverse_end - 1; t >= chunk_start; --t) {
+            const int local_t = t - chunk_start;
+            const float* s_prev = replay_bh +
+                static_cast<size_t>(local_t) * state_elements;
+            const float* s_after = replay_bh +
+                static_cast<size_t>(local_t + 1) * state_elements;
+            const float* q_t = q_bh + t * DR_D_K;
+            const float* k_t = k_bh + t * DR_D_K;
+            const float g_t = g_bh[t];
+            const float beta_t = beta_bh[t];
+            const float go_t = go_bh[t * DR_D_V + tid];
+            const float delta_t =
+                delta_buf[bh * S * DR_D_V + t * DR_D_V + tid];
+
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                grad_s[idx] += q_t[dk] * go_t;
+                float q_part = s_after[idx] * go_t;
+                for (int off = 16; off > 0; off >>= 1)
+                    q_part += __shfl_down_sync(0xffffffff, q_part, off);
+                if (lane == 0)
+                    reduce_q[warp * DR_D_K + dk] = q_part;
+            }
+            __syncthreads();
+            if (tid < DR_D_K) {
+                gq_bh[t * DR_D_K + tid] =
+                    reduce_q[tid] + reduce_q[DR_D_K + tid] +
+                    reduce_q[2 * DR_D_K + tid] +
+                    reduce_q[3 * DR_D_K + tid];
+            }
+
+            float gdelta = 0.0f;
+            float kv_mem = 0.0f;
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                gdelta += grad_s[idx] * k_t[dk];
+                kv_mem += (s_prev[idx] * g_t) * k_t[dk];
+            }
+            const float h = gdelta * beta_t;
+            gv_bh[t * DR_D_V + tid] = h;
+            const float beta_part =
+                gdelta * (v_bh[t * DR_D_V + tid] - kv_mem);
+            float g_part = 0.0f;
+
+            for (int dk = 0; dk < DR_D_K; ++dk) {
+                const int idx = dk * DR_D_V + tid;
+                const float r = s_prev[idx] * g_t;
+                const float grad_r = grad_s[idx] - k_t[dk] * h;
+                const float k_part = grad_s[idx] * delta_t - r * h;
+                g_part += grad_r * s_prev[idx];
+                float reduced_k = k_part;
+                for (int off = 16; off > 0; off >>= 1)
+                    reduced_k += __shfl_down_sync(
+                        0xffffffff, reduced_k, off);
+                if (lane == 0)
+                    reduce_k[warp * DR_D_K + dk] = reduced_k;
+                grad_s[idx] = grad_r * g_t;
+            }
+
+            float reduced_beta = beta_part;
+            float reduced_g = g_part;
+            for (int off = 16; off > 0; off >>= 1) {
+                reduced_beta += __shfl_down_sync(
+                    0xffffffff, reduced_beta, off);
+                reduced_g += __shfl_down_sync(
+                    0xffffffff, reduced_g, off);
+            }
+            if (lane == 0) {
+                reduce_beta[warp] = reduced_beta;
+                reduce_g[warp] = reduced_g;
+            }
+            __syncthreads();
+            if (tid < DR_D_K) {
+                gk_bh[t * DR_D_K + tid] =
+                    reduce_k[tid] + reduce_k[DR_D_K + tid] +
+                    reduce_k[2 * DR_D_K + tid] +
+                    reduce_k[3 * DR_D_K + tid];
+            }
+            if (tid == 0) {
+                gb_bh[t] = reduce_beta[0] + reduce_beta[1] +
+                           reduce_beta[2] + reduce_beta[3];
+                gg_bh[t] = reduce_g[0] + reduce_g[1] +
+                           reduce_g[2] + reduce_g[3];
+            }
+            __syncthreads();
+        }
+    }
+}
+
 // ──────────────────────────────────────────────────────────────────────
 // Host launcher (backward)
 // ──────────────────────────────────────────────────────────────────────
 
-inline void launch_gated_delta_rule_backward(
+inline int launch_gated_delta_rule_backward(
     const float* q, const float* k, const float* v,
     const float* g_exp, const float* beta,
     const float* final_state,
     const float* delta_buf,
+    const float* state_checkpoints, int checkpoint_stride,
     const float* grad_out,
     float* grad_q, float* grad_k, float* grad_v,
     float* grad_g, float* grad_beta,
+    const int32_t* lengths, int heads_per_batch,
     int BH, int seq_len, int key_dim, int val_dim,
-    cudaStream_t stream = 0
+    cudaStream_t stream
 ) {
-    if (key_dim != DR_D_K || val_dim != DR_D_V) {
+    if (key_dim != DR_D_K || val_dim != DR_D_V || heads_per_batch <= 0) {
         fprintf(stderr, "[delta_rule_backward] ERROR: D_K=%d or D_V=%d mismatch\n", key_dim, val_dim);
-        return;
+        return -1;
+    }
+
+    const char* chunkwise_env = std::getenv("QWEN36_GDN_CHUNKWISE_BWD");
+    const bool chunkwise = chunkwise_env && chunkwise_env[0] != '\0' &&
+        std::strcmp(chunkwise_env, "0") != 0 &&
+        std::strcmp(chunkwise_env, "false") != 0;
+    const char* inverse_env = std::getenv("QWEN36_GDN_INVERSE_BWD");
+    const bool inverse = inverse_env && inverse_env[0] != '\0' &&
+        std::strcmp(inverse_env, "0") != 0 &&
+        std::strcmp(inverse_env, "false") != 0;
+
+    // The default path replays each forward chunk from its exact boundary
+    // state. The checkpoint at index zero is required because the state input
+    // is updated in place by forward and final_state cannot recover it stably.
+    // QWEN36_GDN_INVERSE_BWD=1 retains the old low-memory inverse recurrence
+    // solely as an explicit diagnostic/compatibility opt-in.
+    if (!chunkwise && !inverse) {
+        if (state_checkpoints == nullptr || checkpoint_stride <= 0) {
+            fprintf(stderr,
+                    "[delta_rule_backward] stable replay backward requires "
+                    "forward state checkpoints; set a positive checkpoint "
+                    "stride\n");
+            return static_cast<int>(cudaErrorInvalidValue);
+        }
+        if (BH <= 0 || seq_len <= 0) {
+            fprintf(stderr,
+                    "[delta_rule_backward] invalid replay dimensions: "
+                    "BH=%d, S=%d\n", BH, seq_len);
+            return static_cast<int>(cudaErrorInvalidValue);
+        }
+
+        constexpr size_t state_elements = DR_D_K * DR_D_V;
+        const int replay_stride = checkpoint_stride < seq_len
+            ? checkpoint_stride : seq_len;
+        const size_t replay_slots = static_cast<size_t>(replay_stride) + 1;
+        if (static_cast<size_t>(BH) >
+            std::numeric_limits<size_t>::max() / replay_slots /
+                state_elements / sizeof(float)) {
+            fprintf(stderr,
+                    "[delta_rule_backward] replay workspace allocation "
+                    "overflow\n");
+            return static_cast<int>(cudaErrorInvalidValue);
+        }
+        const size_t replay_bytes = static_cast<size_t>(BH) * replay_slots *
+            state_elements * sizeof(float);
+        float* replay_states = nullptr;
+        auto alloc_status = cudaMallocAsync(
+            reinterpret_cast<void**>(&replay_states), replay_bytes, stream);
+        if (alloc_status != cudaSuccess) {
+            fprintf(stderr,
+                    "[delta_rule_backward] replay workspace allocation "
+                    "failed (%zu bytes): %s\n", replay_bytes,
+                    cudaGetErrorString(alloc_status));
+            return static_cast<int>(alloc_status);
+        }
+
+        const size_t replay_smem =
+            (2 * state_elements + 8 * DR_D_K + 8) * sizeof(float);
+        auto attr_status = cudaFuncSetAttribute(
+            gated_delta_rule_backward_replay_kernel,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, replay_smem);
+        if (attr_status != cudaSuccess) {
+            cudaFreeAsync(replay_states, stream);
+            fprintf(stderr,
+                    "[delta_rule_backward] replay shared-memory attribute "
+                    "failed: %s\n", cudaGetErrorString(attr_status));
+            return static_cast<int>(attr_status);
+        }
+        gated_delta_rule_backward_replay_kernel
+            <<<dim3(BH), dim3(DR_THREADS), replay_smem, stream>>>(
+                q, k, v, g_exp, beta, delta_buf, state_checkpoints,
+                replay_stride, replay_states, grad_out,
+                grad_q, grad_k, grad_v, grad_g, grad_beta,
+                lengths, heads_per_batch, seq_len);
+        const auto launch_status = cudaGetLastError();
+        const auto free_status = cudaFreeAsync(replay_states, stream);
+        if (launch_status != cudaSuccess)
+            return static_cast<int>(launch_status);
+        return free_status == cudaSuccess ? 0 : static_cast<int>(free_status);
+    }
+
+    if (chunkwise && (state_checkpoints == nullptr || checkpoint_stride <= 0)) {
+        fprintf(stderr,
+                "[delta_rule_backward] chunkwise backward requires a state "
+                "checkpoint stride no greater than the sequence length\n");
+        return static_cast<int>(cudaErrorInvalidValue);
+    }
+    if (chunkwise) {
+        const int checkpoint_count =
+            (seq_len + checkpoint_stride - 1) / checkpoint_stride + 1;
+        const int num_chunks = checkpoint_count - 1;
+        constexpr size_t state_elements = DR_D_K * DR_D_V;
+        if (static_cast<size_t>(BH) >
+            std::numeric_limits<size_t>::max() /
+                static_cast<size_t>(checkpoint_count) /
+                state_elements / sizeof(float)) {
+            fprintf(stderr,
+                    "[delta_rule_backward] chunk boundary allocation overflow\n");
+            return static_cast<int>(cudaErrorInvalidValue);
+        }
+        const size_t boundary_bytes = static_cast<size_t>(BH) *
+            checkpoint_count * state_elements * sizeof(float);
+        float* grad_boundaries = nullptr;
+        auto alloc_status = cudaMallocAsync(
+            reinterpret_cast<void**>(&grad_boundaries), boundary_bytes, stream);
+        if (alloc_status != cudaSuccess) {
+            fprintf(stderr,
+                    "[delta_rule_backward] chunk boundary allocation failed: %s\n",
+                    cudaGetErrorString(alloc_status));
+            return static_cast<int>(alloc_status);
+        }
+
+        const size_t boundary_smem = state_elements * sizeof(float);
+        auto attr_status = cudaFuncSetAttribute(
+            gated_delta_rule_backward_boundaries_kernel,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, boundary_smem);
+        if (attr_status != cudaSuccess) {
+            cudaFreeAsync(grad_boundaries, stream);
+            fprintf(stderr,
+                    "[delta_rule_backward] boundary shared-memory attribute failed: %s\n",
+                    cudaGetErrorString(attr_status));
+            return static_cast<int>(attr_status);
+        }
+        gated_delta_rule_backward_boundaries_kernel
+            <<<dim3(BH), dim3(DR_THREADS), boundary_smem, stream>>>(
+                q, k, g_exp, beta, grad_out, grad_boundaries,
+                lengths, heads_per_batch, seq_len, checkpoint_stride,
+                checkpoint_count);
+        auto launch_status = cudaGetLastError();
+        if (launch_status != cudaSuccess) {
+            cudaFreeAsync(grad_boundaries, stream);
+            return static_cast<int>(launch_status);
+        }
+
+        const size_t chunk_smem =
+            (2 * state_elements + 8 * DR_D_K + 8) * sizeof(float);
+        attr_status = cudaFuncSetAttribute(
+            gated_delta_rule_backward_chunks_kernel,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, chunk_smem);
+        if (attr_status != cudaSuccess) {
+            cudaFreeAsync(grad_boundaries, stream);
+            fprintf(stderr,
+                    "[delta_rule_backward] chunk shared-memory attribute failed: %s\n",
+                    cudaGetErrorString(attr_status));
+            return static_cast<int>(attr_status);
+        }
+        gated_delta_rule_backward_chunks_kernel
+            <<<dim3(BH * num_chunks), dim3(DR_THREADS), chunk_smem, stream>>>(
+                q, k, v, g_exp, beta, delta_buf, state_checkpoints,
+                grad_boundaries, checkpoint_stride, checkpoint_count, grad_out,
+                grad_q, grad_k, grad_v, grad_g, grad_beta,
+                lengths, heads_per_batch, seq_len);
+        launch_status = cudaGetLastError();
+        const auto free_status = cudaFreeAsync(grad_boundaries, stream);
+        if (launch_status != cudaSuccess)
+            return static_cast<int>(launch_status);
+        return free_status == cudaSuccess ? 0 : static_cast<int>(free_status);
     }
 
+    // Once inverse recurrence is explicitly enabled, this flag selects its
+    // fused or checkpoint-reload implementation. Neither is a default path.
+    const char* fusion_env = std::getenv("QWEN36_GDN_RECURRENT_FUSION");
+    const bool fuse_recurrence = !fusion_env || fusion_env[0] == '\0' ||
+        (std::strcmp(fusion_env, "0") != 0 &&
+         std::strcmp(fusion_env, "false") != 0);
     dim3 grid(BH);
     dim3 block(DR_THREADS);
 
-    // 2 state matrices in shared memory: 2 * 64KB = 128KB
-    size_t smem_size = 2 * DR_D_K * DR_D_V * sizeof(float);
-    cudaFuncSetAttribute(
-        gated_delta_rule_backward_kernel,
-        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+    // 2 state matrices plus four warp partials per D_K and scalar reductions.
+    size_t smem_size = (2 * DR_D_K * DR_D_V + 8 * DR_D_K + 8) * sizeof(float);
+    const void* kernel = fuse_recurrence
+        ? reinterpret_cast<const void*>(
+              gated_delta_rule_backward_kernel_correct<true>)
+        : reinterpret_cast<const void*>(
+              gated_delta_rule_backward_kernel_correct<false>);
+    auto attr_status = cudaFuncSetAttribute(
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+    if (attr_status != cudaSuccess) {
+        fprintf(stderr, "[delta_rule_backward] shared-memory attribute failed: %s\n",
+                cudaGetErrorString(attr_status));
+        return static_cast<int>(attr_status);
+    }
 
-    gated_delta_rule_backward_kernel<<<grid, block, smem_size, stream>>>(
-        q, k, v, g_exp, beta, final_state, delta_buf, grad_out,
-        grad_q, grad_k, grad_v, grad_g, grad_beta, seq_len
-    );
+    if (fuse_recurrence) {
+        gated_delta_rule_backward_kernel_correct<true>
+            <<<grid, block, smem_size, stream>>>(
+                q, k, v, g_exp, beta, final_state, delta_buf,
+                state_checkpoints, checkpoint_stride, grad_out,
+                grad_q, grad_k, grad_v, grad_g, grad_beta,
+                lengths, heads_per_batch, seq_len);
+    } else {
+        gated_delta_rule_backward_kernel_correct<false>
+            <<<grid, block, smem_size, stream>>>(
+                q, k, v, g_exp, beta, final_state, delta_buf,
+                state_checkpoints, checkpoint_stride, grad_out,
+                grad_q, grad_k, grad_v, grad_g, grad_beta,
+                lengths, heads_per_batch, seq_len);
+    }
+    return 0;
 }
diff --git a/crates/rustrain-qwen3-6/kernels/fused_kernels.cu b/crates/rustrain-qwen3-6/kernels/fused_kernels.cu
index c3894d05..45288bd3 100644
--- a/crates/rustrain-qwen3-6/kernels/fused_kernels.cu
+++ b/crates/rustrain-qwen3-6/kernels/fused_kernels.cu
@@ -10,6 +10,8 @@
 
 #include <cuda_runtime.h>
 #include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cstdint>
 #include <cmath>
 
 // ──────────────────────────────────────────────────────────────────────
@@ -184,10 +186,223 @@ __global__ void fused_rmsnorm_matmul_kernel(
 }
 
 // ──────────────────────────────────────────────────────────────────────
-// 4. Multi-tensor Fused Adam
+// 4. Fused MoE weighted unpermute
+// ──────────────────────────────────────────────────────────────────────
+
+template <typename T>
+__device__ __forceinline__ float moe_to_float(T value);
+
+template <>
+__device__ __forceinline__ float moe_to_float<float>(float value) {
+    return value;
+}
+
+template <>
+__device__ __forceinline__ float moe_to_float<__nv_bfloat16>(
+    __nv_bfloat16 value
+) {
+    return __bfloat162float(value);
+}
+
+template <>
+__device__ __forceinline__ float moe_to_float<__half>(__half value) {
+    return __half2float(value);
+}
+
+template <typename T>
+__device__ __forceinline__ T moe_from_float(float value);
+
+template <>
+__device__ __forceinline__ float moe_from_float<float>(float value) {
+    return value;
+}
+
+template <>
+__device__ __forceinline__ __nv_bfloat16 moe_from_float<__nv_bfloat16>(
+    float value
+) {
+    return __float2bfloat16_rn(value);
+}
+
+template <>
+__device__ __forceinline__ __half moe_from_float<__half>(float value) {
+    return __float2half_rn(value);
+}
+
+__global__ void moe_inverse_order_kernel(
+    const int64_t* __restrict__ send_index,
+    int64_t* __restrict__ inverse_order,
+    int64_t assignments
+) {
+    const int64_t sorted =
+        static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (sorted >= assignments) return;
+    inverse_order[send_index[sorted]] = sorted;
+}
+
+template <typename T>
+__global__ void moe_weighted_unpermute_forward_kernel(
+    const T* __restrict__ returned,
+    const T* __restrict__ assignment_weights,
+    const int64_t* __restrict__ inverse_order,
+    T* __restrict__ output,
+    int64_t tokens,
+    int64_t hidden,
+    int top_k
+) {
+    const int64_t token = blockIdx.x;
+    if (token >= tokens) return;
+    for (int64_t column = threadIdx.x;
+         column < hidden; column += blockDim.x) {
+        float sum = 0.0f;
+        const int64_t first_assignment = token * top_k;
+        for (int route = 0; route < top_k; ++route) {
+            const int64_t original = first_assignment + route;
+            const int64_t sorted = inverse_order[original];
+            sum += moe_to_float(returned[sorted * hidden + column]) *
+                moe_to_float(assignment_weights[original]);
+        }
+        output[token * hidden + column] = moe_from_float<T>(sum);
+    }
+}
+
+template <typename T>
+__global__ void moe_weighted_unpermute_backward_kernel(
+    const T* __restrict__ returned,
+    const T* __restrict__ assignment_weights,
+    const int64_t* __restrict__ inverse_order,
+    const T* __restrict__ grad_output,
+    T* __restrict__ grad_returned,
+    T* __restrict__ grad_assignment_weights,
+    int64_t assignments,
+    int64_t hidden,
+    int top_k
+) {
+    const int64_t original = blockIdx.x;
+    if (original >= assignments) return;
+    const int64_t token = original / top_k;
+    const int64_t sorted = inverse_order[original];
+    const float weight = moe_to_float(assignment_weights[original]);
+    float weight_grad = 0.0f;
+    for (int64_t column = threadIdx.x;
+         column < hidden; column += blockDim.x) {
+        const float grad = moe_to_float(grad_output[token * hidden + column]);
+        weight_grad += grad * moe_to_float(returned[sorted * hidden + column]);
+        grad_returned[sorted * hidden + column] =
+            moe_from_float<T>(grad * weight);
+    }
+
+    for (int offset = 16; offset > 0; offset >>= 1)
+        weight_grad += __shfl_down_sync(0xffffffff, weight_grad, offset);
+    __shared__ float warp_sums[8];
+    const int lane = threadIdx.x & 31;
+    const int warp = threadIdx.x >> 5;
+    if (lane == 0) warp_sums[warp] = weight_grad;
+    __syncthreads();
+    if (warp == 0) {
+        weight_grad = lane < (blockDim.x + 31) / 32
+            ? warp_sums[lane] : 0.0f;
+        for (int offset = 16; offset > 0; offset >>= 1)
+            weight_grad += __shfl_down_sync(
+                0xffffffff, weight_grad, offset);
+        if (lane == 0)
+            grad_assignment_weights[original] =
+                moe_from_float<T>(weight_grad);
+    }
+}
+
+// GPU counting sort for the receive-side local-expert permutation. The
+// sorted_to_received mapping is the single source of truth for activation and
+// token metadata alignment, and is reused to invert the permutation.
+__global__ void moe_local_expert_histogram_kernel(
+    const int64_t* __restrict__ received_experts,
+    int64_t expert_stride,
+    int* __restrict__ counts,
+    int64_t rows,
+    int expert_count
+) {
+    const int64_t row =
+        static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (row >= rows) return;
+    const int64_t expert = received_experts[row * expert_stride];
+    if (expert >= 0 && expert < expert_count)
+        atomicAdd(counts + expert, 1);
+}
+
+__global__ void moe_local_expert_prefix_kernel(
+    const int* __restrict__ counts,
+    int* __restrict__ offsets,
+    int* __restrict__ cursors,
+    int expert_count
+) {
+    __shared__ int scan[1024];
+    const int tid = threadIdx.x;
+    if (tid < expert_count) scan[tid] = counts[tid];
+
+    for (int distance = 1; distance < expert_count; distance <<= 1) {
+        __syncthreads();
+        const int value = tid < expert_count ? scan[tid] : 0;
+        const int prefix = tid >= distance && tid < expert_count
+            ? scan[tid - distance] : 0;
+        __syncthreads();
+        if (tid < expert_count) scan[tid] = value + prefix;
+    }
+    __syncthreads();
+    if (tid < expert_count) {
+        offsets[tid] = scan[tid];
+        cursors[tid] = tid == 0 ? 0 : scan[tid - 1];
+    }
+}
+
+__global__ void moe_local_expert_scatter_metadata_kernel(
+    const int64_t* __restrict__ received_tokens,
+    int64_t token_stride,
+    const int64_t* __restrict__ received_experts,
+    int64_t expert_stride,
+    int* __restrict__ cursors,
+    int64_t* __restrict__ selected_tokens,
+    int64_t* __restrict__ sorted_experts,
+    int64_t* __restrict__ sorted_to_received,
+    int64_t rows,
+    int expert_count
+) {
+    const int64_t received_row =
+        static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (received_row >= rows) return;
+    const int64_t expert = received_experts[received_row * expert_stride];
+    if (expert < 0 || expert >= expert_count) return;
+    const int sorted_row = atomicAdd(cursors + expert, 1);
+    sorted_to_received[sorted_row] = received_row;
+    selected_tokens[sorted_row] =
+        received_tokens[received_row * token_stride];
+    sorted_experts[sorted_row] = expert;
+}
+
+template <typename T, bool Inverse>
+__global__ void moe_local_permute_rows_kernel(
+    const T* __restrict__ input,
+    const int64_t* __restrict__ sorted_to_received,
+    T* __restrict__ output,
+    int64_t rows,
+    int64_t hidden
+) {
+    const int64_t sorted_row = blockIdx.x;
+    if (sorted_row >= rows) return;
+    const int64_t received_row = sorted_to_received[sorted_row];
+    const int64_t source_row = Inverse ? sorted_row : received_row;
+    const int64_t destination_row = Inverse ? received_row : sorted_row;
+    for (int64_t column = threadIdx.x;
+         column < hidden; column += blockDim.x) {
+        output[destination_row * hidden + column] =
+            input[source_row * hidden + column];
+    }
+}
+
+// ──────────────────────────────────────────────────────────────────────
+// 5. Multi-tensor Fused Adam
 // ──────────────────────────────────────────────────────────────────────
 // One block per param tensor. Each thread processes multiple elements.
-// Handles BF16 params/grads with FP32 m/v.
+// Handles BF16 params with FP32 accumulated grads and FP32 m/v.
 //
 // Replaces 7 ATen ops per param with 1 kernel launch for ALL params:
 //   m = m * beta1 + grad * (1 - beta1)
@@ -197,7 +412,7 @@ __global__ void fused_rmsnorm_matmul_kernel(
 
 __global__ void fused_adam_multi_kernel(
     void** __restrict__ param_ptrs,   // [n_params] BF16
-    void** __restrict__ grad_ptrs,    // [n_params] BF16
+    void** __restrict__ grad_ptrs,    // [n_params] FP32
     float** __restrict__ m_ptrs,      // [n_params] FP32
     float** __restrict__ v_ptrs,      // [n_params] FP32
     const int* __restrict__ sizes,   // [n_params]
@@ -211,12 +426,12 @@ __global__ void fused_adam_multi_kernel(
 
     int size = sizes[pidx];
     __nv_bfloat16* param = (__nv_bfloat16*)param_ptrs[pidx];
-    __nv_bfloat16* grad  = (__nv_bfloat16*)grad_ptrs[pidx];
+    float* grad = (float*)grad_ptrs[pidx];
     float* m = m_ptrs[pidx];
     float* v = v_ptrs[pidx];
 
     for (int i = threadIdx.x; i < size; i += blockDim.x) {
-        float g = __bfloat162float(grad[i]);
+        float g = grad[i];
         float m_new = m[i] * beta1 + g * one_minus_beta1;
         float v_new = v[i] * beta2 + g * g * one_minus_beta2;
         m[i] = m_new;
@@ -224,10 +439,117 @@ __global__ void fused_adam_multi_kernel(
         float p = __bfloat162float(param[i]);
         p -= lr_scaled * m_new / (sqrtf(v_new) + eps_scaled);
         param[i] = __float2bfloat16_rn(p);
-        grad[i] = __float2bfloat16_rn(0.0f);
+        grad[i] = 0.0f;
+    }
+}
+
+// Out-of-place variant used by dynamic multi-LoRA transactions. It writes
+// only destination parameter/state tensors and never mutates source tensors
+// or accumulated gradients.
+__global__ void fused_adam_multi_out_of_place_kernel(
+    void** __restrict__ src_param_ptrs,
+    void** __restrict__ grad_ptrs,
+    float** __restrict__ src_m_ptrs,
+    float** __restrict__ src_v_ptrs,
+    void** __restrict__ dst_param_ptrs,
+    float** __restrict__ dst_m_ptrs,
+    float** __restrict__ dst_v_ptrs,
+    const int* __restrict__ sizes,
+    const float* __restrict__ lr_scaled,
+    const float* __restrict__ eps_scaled,
+    const float* __restrict__ beta1,
+    const float* __restrict__ beta2,
+    int n_params
+) {
+    int pidx = blockIdx.x;
+    if (pidx >= n_params) return;
+
+    int size = sizes[pidx];
+    const __nv_bfloat16* src_param =
+        (const __nv_bfloat16*)src_param_ptrs[pidx];
+    const float* grad = (const float*)grad_ptrs[pidx];
+    const float* src_m = src_m_ptrs[pidx];
+    const float* src_v = src_v_ptrs[pidx];
+    __nv_bfloat16* dst_param = (__nv_bfloat16*)dst_param_ptrs[pidx];
+    float* dst_m = dst_m_ptrs[pidx];
+    float* dst_v = dst_v_ptrs[pidx];
+    const float tensor_lr = lr_scaled[pidx];
+    const float tensor_eps = eps_scaled[pidx];
+    const float tensor_beta1 = beta1[pidx];
+    const float tensor_beta2 = beta2[pidx];
+    const float tensor_one_minus_beta1 = 1.0f - tensor_beta1;
+    const float tensor_one_minus_beta2 = 1.0f - tensor_beta2;
+
+    for (int i = threadIdx.x; i < size; i += blockDim.x) {
+        float g = grad[i];
+        float m_new = src_m[i] * tensor_beta1 +
+            g * tensor_one_minus_beta1;
+        float v_new = src_v[i] * tensor_beta2 +
+            g * g * tensor_one_minus_beta2;
+        dst_m[i] = m_new;
+        dst_v[i] = v_new;
+        float p = __bfloat162float(src_param[i]);
+        p -= tensor_lr * m_new / (sqrtf(v_new) + tensor_eps);
+        dst_param[i] = __float2bfloat16_rn(p);
+    }
+}
+
+// Accumulate one FP32 squared L2 norm per logical adapter. Pointer lists only
+// contain the unique owners of replicated parameters, so the caller can sum
+// these device scalars over the orthogonal process grid without double count.
+__global__ void fused_multi_tensor_l2_norm_kernel(
+    void** __restrict__ grad_ptrs,
+    const int* __restrict__ sizes,
+    const int* __restrict__ groups,
+    int n_tensors,
+    float* __restrict__ norm_squares
+) {
+    const int tensor_index = blockIdx.x;
+    if (tensor_index >= n_tensors) return;
+
+    const float* grad = (const float*)grad_ptrs[tensor_index];
+    const int size = sizes[tensor_index];
+    float sum = 0.0f;
+    for (int i = threadIdx.x; i < size; i += blockDim.x) {
+        const float value = grad[i];
+        sum += value * value;
+    }
+    for (int offset = 16; offset > 0; offset >>= 1)
+        sum += __shfl_down_sync(0xffffffff, sum, offset);
+
+    __shared__ float warp_sums[8];
+    const int lane = threadIdx.x & 31;
+    const int warp = threadIdx.x >> 5;
+    if (lane == 0) warp_sums[warp] = sum;
+    __syncthreads();
+    if (warp == 0) {
+        sum = lane < (blockDim.x + 31) / 32 ? warp_sums[lane] : 0.0f;
+        for (int offset = 16; offset > 0; offset >>= 1)
+            sum += __shfl_down_sync(0xffffffff, sum, offset);
+        if (lane == 0) atomicAdd(norm_squares + groups[tensor_index], sum);
     }
 }
 
+__global__ void fused_multi_tensor_clip_kernel(
+    void** __restrict__ grad_ptrs,
+    const int* __restrict__ sizes,
+    const int* __restrict__ groups,
+    int n_tensors,
+    const float* __restrict__ norm_squares,
+    float max_norm
+) {
+    const int tensor_index = blockIdx.x;
+    if (tensor_index >= n_tensors) return;
+    const float total_norm = sqrtf(norm_squares[groups[tensor_index]]);
+    const float scale = fminf(1.0f, max_norm / (total_norm + 1.0e-6f));
+    if (scale >= 1.0f) return;
+
+    float* grad = (float*)grad_ptrs[tensor_index];
+    const int size = sizes[tensor_index];
+    for (int i = threadIdx.x; i < size; i += blockDim.x)
+        grad[i] *= scale;
+}
+
 // ──────────────────────────────────────────────────────────────────────
 // C wrapper functions (called from C++ via extern "C")
 // ──────────────────────────────────────────────────────────────────────
@@ -273,6 +595,194 @@ void launch_fused_rmsnorm_matmul(
     );
 }
 
+void launch_fused_moe_weighted_unpermute_forward(
+    const void* returned, const void* assignment_weights,
+    const int64_t* send_index, int64_t* inverse_order, void* output,
+    int64_t assignments, int64_t tokens, int64_t hidden, int top_k,
+    int dtype, cudaStream_t stream
+) {
+    if (assignments <= 0) return;
+    constexpr int threads = 256;
+    const int inverse_blocks = static_cast<int>(
+        (assignments + threads - 1) / threads);
+    moe_inverse_order_kernel<<<inverse_blocks, threads, 0, stream>>>(
+        send_index, inverse_order, assignments);
+    switch (dtype) {
+        case 0:
+            moe_weighted_unpermute_forward_kernel<<<
+                static_cast<unsigned int>(tokens), threads, 0, stream>>>(
+                static_cast<const __nv_bfloat16*>(returned),
+                static_cast<const __nv_bfloat16*>(assignment_weights),
+                inverse_order, static_cast<__nv_bfloat16*>(output),
+                tokens, hidden, top_k);
+            break;
+        case 1:
+            moe_weighted_unpermute_forward_kernel<<<
+                static_cast<unsigned int>(tokens), threads, 0, stream>>>(
+                static_cast<const float*>(returned),
+                static_cast<const float*>(assignment_weights),
+                inverse_order, static_cast<float*>(output),
+                tokens, hidden, top_k);
+            break;
+        case 2:
+            moe_weighted_unpermute_forward_kernel<<<
+                static_cast<unsigned int>(tokens), threads, 0, stream>>>(
+                static_cast<const __half*>(returned),
+                static_cast<const __half*>(assignment_weights),
+                inverse_order, static_cast<__half*>(output),
+                tokens, hidden, top_k);
+            break;
+    }
+}
+
+void launch_fused_moe_weighted_unpermute_backward(
+    const void* returned, const void* assignment_weights,
+    const int64_t* inverse_order, const void* grad_output,
+    void* grad_returned, void* grad_assignment_weights,
+    int64_t assignments, int64_t hidden, int top_k,
+    int dtype, cudaStream_t stream
+) {
+    if (assignments <= 0) return;
+    constexpr int threads = 256;
+    const auto blocks = static_cast<unsigned int>(assignments);
+    switch (dtype) {
+        case 0:
+            moe_weighted_unpermute_backward_kernel<<<blocks, threads, 0, stream>>>(
+                static_cast<const __nv_bfloat16*>(returned),
+                static_cast<const __nv_bfloat16*>(assignment_weights),
+                inverse_order, static_cast<const __nv_bfloat16*>(grad_output),
+                static_cast<__nv_bfloat16*>(grad_returned),
+                static_cast<__nv_bfloat16*>(grad_assignment_weights),
+                assignments, hidden, top_k);
+            break;
+        case 1:
+            moe_weighted_unpermute_backward_kernel<<<blocks, threads, 0, stream>>>(
+                static_cast<const float*>(returned),
+                static_cast<const float*>(assignment_weights),
+                inverse_order, static_cast<const float*>(grad_output),
+                static_cast<float*>(grad_returned),
+                static_cast<float*>(grad_assignment_weights),
+                assignments, hidden, top_k);
+            break;
+        case 2:
+            moe_weighted_unpermute_backward_kernel<<<blocks, threads, 0, stream>>>(
+                static_cast<const __half*>(returned),
+                static_cast<const __half*>(assignment_weights),
+                inverse_order, static_cast<const __half*>(grad_output),
+                static_cast<__half*>(grad_returned),
+                static_cast<__half*>(grad_assignment_weights),
+                assignments, hidden, top_k);
+            break;
+    }
+}
+
+void launch_fused_moe_local_permute_forward(
+    const void* received,
+    const int64_t* received_tokens, int64_t token_stride,
+    const int64_t* received_experts, int64_t expert_stride,
+    void* selected, int64_t* selected_tokens, int64_t* sorted_experts,
+    int* counts, int* offsets, int* cursors,
+    int64_t* sorted_to_received,
+    int64_t rows, int64_t hidden, int expert_count,
+    int dtype, cudaStream_t stream
+) {
+    constexpr int threads = 256;
+    cudaMemsetAsync(counts, 0, expert_count * sizeof(int), stream);
+    if (rows > 0) {
+        const int blocks = static_cast<int>((rows + threads - 1) / threads);
+        moe_local_expert_histogram_kernel<<<blocks, threads, 0, stream>>>(
+            received_experts, expert_stride, counts, rows, expert_count);
+    }
+    moe_local_expert_prefix_kernel<<<1, 1024, 0, stream>>>(
+        counts, offsets, cursors, expert_count);
+    if (rows <= 0) return;
+
+    const int metadata_blocks = static_cast<int>(
+        (rows + threads - 1) / threads);
+    moe_local_expert_scatter_metadata_kernel<<<
+        metadata_blocks, threads, 0, stream>>>(
+        received_tokens, token_stride, received_experts, expert_stride,
+        cursors, selected_tokens, sorted_experts, sorted_to_received,
+        rows, expert_count);
+    const auto row_blocks = static_cast<unsigned int>(rows);
+    switch (dtype) {
+        case 0:
+            moe_local_permute_rows_kernel<__nv_bfloat16, false><<<
+                row_blocks, threads, 0, stream>>>(
+                static_cast<const __nv_bfloat16*>(received),
+                sorted_to_received, static_cast<__nv_bfloat16*>(selected),
+                rows, hidden);
+            break;
+        case 1:
+            moe_local_permute_rows_kernel<float, false><<<
+                row_blocks, threads, 0, stream>>>(
+                static_cast<const float*>(received), sorted_to_received,
+                static_cast<float*>(selected), rows, hidden);
+            break;
+        case 2:
+            moe_local_permute_rows_kernel<__half, false><<<
+                row_blocks, threads, 0, stream>>>(
+                static_cast<const __half*>(received), sorted_to_received,
+                static_cast<__half*>(selected), rows, hidden);
+            break;
+    }
+}
+
+void launch_fused_moe_local_permute_rows(
+    const void* input, const int64_t* sorted_to_received, void* output,
+    int64_t rows, int64_t hidden, int dtype, bool inverse,
+    cudaStream_t stream
+) {
+    if (rows <= 0) return;
+    constexpr int threads = 256;
+    const auto blocks = static_cast<unsigned int>(rows);
+    if (inverse) {
+        switch (dtype) {
+            case 0:
+                moe_local_permute_rows_kernel<__nv_bfloat16, true><<<
+                    blocks, threads, 0, stream>>>(
+                    static_cast<const __nv_bfloat16*>(input),
+                    sorted_to_received, static_cast<__nv_bfloat16*>(output),
+                    rows, hidden);
+                break;
+            case 1:
+                moe_local_permute_rows_kernel<float, true><<<
+                    blocks, threads, 0, stream>>>(
+                    static_cast<const float*>(input), sorted_to_received,
+                    static_cast<float*>(output), rows, hidden);
+                break;
+            case 2:
+                moe_local_permute_rows_kernel<__half, true><<<
+                    blocks, threads, 0, stream>>>(
+                    static_cast<const __half*>(input), sorted_to_received,
+                    static_cast<__half*>(output), rows, hidden);
+                break;
+        }
+    } else {
+        switch (dtype) {
+            case 0:
+                moe_local_permute_rows_kernel<__nv_bfloat16, false><<<
+                    blocks, threads, 0, stream>>>(
+                    static_cast<const __nv_bfloat16*>(input),
+                    sorted_to_received, static_cast<__nv_bfloat16*>(output),
+                    rows, hidden);
+                break;
+            case 1:
+                moe_local_permute_rows_kernel<float, false><<<
+                    blocks, threads, 0, stream>>>(
+                    static_cast<const float*>(input), sorted_to_received,
+                    static_cast<float*>(output), rows, hidden);
+                break;
+            case 2:
+                moe_local_permute_rows_kernel<__half, false><<<
+                    blocks, threads, 0, stream>>>(
+                    static_cast<const __half*>(input), sorted_to_received,
+                    static_cast<__half*>(output), rows, hidden);
+                break;
+        }
+    }
+}
+
 void launch_fused_adam_multi(
     void** d_param_ptrs, void** d_grad_ptrs,
     float** d_m_ptrs, float** d_v_ptrs,
@@ -293,4 +803,46 @@ void launch_fused_adam_multi(
     );
 }
 
+void launch_fused_adam_multi_out_of_place(
+    void** d_src_param_ptrs, void** d_grad_ptrs,
+    float** d_src_m_ptrs, float** d_src_v_ptrs,
+    void** d_dst_param_ptrs,
+    float** d_dst_m_ptrs, float** d_dst_v_ptrs,
+    int* d_sizes, float* d_lr_scaled, float* d_eps_scaled,
+    float* d_beta1, float* d_beta2,
+    int n_params,
+    cudaStream_t stream
+) {
+    if (n_params <= 0) return;
+    int threads = 256;
+    int blocks = n_params;
+    fused_adam_multi_out_of_place_kernel<<<blocks, threads, 0, stream>>>(
+        d_src_param_ptrs, d_grad_ptrs, d_src_m_ptrs, d_src_v_ptrs,
+        d_dst_param_ptrs, d_dst_m_ptrs, d_dst_v_ptrs,
+        d_sizes, d_lr_scaled, d_eps_scaled,
+        d_beta1, d_beta2,
+        n_params
+    );
+}
+
+void launch_fused_multi_tensor_l2_norm(
+    void** d_grad_ptrs, int* d_sizes, int* d_groups,
+    int n_tensors, float* d_norm_squares, cudaStream_t stream
+) {
+    if (n_tensors <= 0) return;
+    fused_multi_tensor_l2_norm_kernel<<<n_tensors, 256, 0, stream>>>(
+        d_grad_ptrs, d_sizes, d_groups, n_tensors, d_norm_squares);
+}
+
+void launch_fused_multi_tensor_clip(
+    void** d_grad_ptrs, int* d_sizes, int* d_groups,
+    int n_tensors, const float* d_norm_squares, float max_norm,
+    cudaStream_t stream
+) {
+    if (n_tensors <= 0) return;
+    fused_multi_tensor_clip_kernel<<<n_tensors, 256, 0, stream>>>(
+        d_grad_ptrs, d_sizes, d_groups, n_tensors,
+        d_norm_squares, max_norm);
+}
+
 }  // extern "C"
diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp
index ba76c46e..795439c2 100644
--- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp
+++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp
@@ -5,7 +5,14 @@
 // No tch-rs VarStore involved — gradients flow entirely within C++ autograd.
 
 #include <ATen/ATen.h>
+#if __has_include(<ATen/ops/_grouped_mm.h>)
+#include <ATen/ops/_grouped_mm.h>
+#define RUSTRAIN_HAS_ATEN_GROUPED_MM 1
+#else
+#define RUSTRAIN_HAS_ATEN_GROUPED_MM 0
+#endif
 #include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDACachingAllocator.h>
 #include <cuda_runtime.h>
 #include <torch/csrc/autograd/grad_mode.h>
@@ -13,7 +20,11 @@
 #include <torch/csrc/autograd/autograd.h>
 #include <torch/csrc/autograd/variable.h>
 #include <cstdio>
+#include <cstdlib>
+#include <cstdint>
 #include <cmath>
+#include <algorithm>
+#include <iterator>
 #include <vector>
 #include <cstring>
 #include <memory>
@@ -21,8 +32,10 @@
 #include <set>
 #include <array>
 #include <map>
+#include <limits>
 #include <sstream>
 #include <chrono>
+#include <cerrno>
 #include <sys/stat.h>
 #include <unistd.h>
 
@@ -30,10 +43,73 @@
 #include <nccl.h>
 
 struct TrainingContext;  // forward declaration (defined below)
+static cudaStream_t context_moe_shared_stream(TrainingContext* ctx);
+static cudaStream_t context_moe_ep_metadata_stream(TrainingContext* ctx);
+
+static at::Tensor attach_router_aux_loss(
+    TrainingContext* ctx,
+    const at::Tensor& routing_weights,
+    const at::Tensor& topk_indices,
+    int64_t batch,
+    int64_t seq,
+    int64_t top_k,
+    int64_t num_experts);
 struct LayerConfig;
+static int64_t g_context_sequence = 0;
+
+// ABI-stable non-interleaved pipeline window contract. The window supports
+// fixed-shape, one-chunk 1F1B execution for any PP size >= 2; the legacy
+// one-microstep entry point remains available as a compatibility fallback.
+struct Qwen36PipelineWindowV1 {
+    uint32_t struct_size;
+    uint32_t version;
+    int64_t window_id;
+    int64_t num_microbatches;
+    int32_t schedule;
+    int32_t num_chunks;
+    int32_t flags;
+};
+
+// Window flags are additive to the v1 ABI. Bit 0 selects the dynamic
+// multi-tenant LoRA finalizer while retaining the fixed 1F1B schedule.
+static constexpr int32_t kPipelineWindowFlagDynamicLora = 1 << 0;
+
+struct Qwen36PipelineTickV1 {
+    uint32_t struct_size;
+    uint32_t version;
+    int64_t window_id;
+    int64_t forward_mb;
+    int64_t backward_mb;
+    int32_t chunk_id;
+    int32_t phase;
+    void* input_ids;
+    void* target_mask;
+    void* attention_mask;
+    double gradient_scale;
+};
+
+struct Qwen36PipelineResultV1 {
+    uint32_t struct_size;
+    uint32_t version;
+    int32_t status;
+    int64_t completed_fwd;
+    int64_t completed_bwd;
+    int64_t in_flight;
+    int64_t optimizer_step;
+    double loss;
+};
+
+enum class DynamicMultiLoraMode;
+extern "C" double qwen36_train_multi_lora_impl(
+    void*, void*, void*, void*, int32_t, int32_t, DynamicMultiLoraMode,
+    int32_t*, at::Tensor*, at::Tensor*, bool, const at::Tensor*, bool, bool);
 
 // Forward declarations for functions defined after TrainingContext
 at::Tensor apply_multi_lora(TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx, const at::Tensor& base_weight);
+static at::Tensor tp_allreduce_lora_delta(
+    TrainingContext* ctx, const at::Tensor& local_delta);
+static at::Tensor tp_copy_lora_input(
+    TrainingContext* ctx, const at::Tensor& input);
 
 // ──────────────────────────────────────────────────────────────────────
 // Forward declarations
@@ -41,6 +117,101 @@ at::Tensor apply_multi_lora(TrainingContext* ctx, int64_t layer_idx, int64_t pai
 
 static at::Tensor rms_norm(const at::Tensor& input, const at::Tensor& weight, double eps);
 
+static bool env_enabled(const char* name, bool fallback = false) {
+    const char* value = std::getenv(name);
+    if (!value || value[0] == '\0') return fallback;
+    return std::strcmp(value, "0") != 0 &&
+        std::strcmp(value, "false") != 0;
+}
+
+static bool valid_rendezvous_component(const char* value) {
+    if (!value || value[0] == '\0' || std::strcmp(value, ".") == 0 ||
+        std::strcmp(value, "..") == 0) {
+        return false;
+    }
+    for (const unsigned char ch : std::string(value)) {
+        if (!((ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') ||
+              (ch >= '0' && ch <= '9') || ch == '-' || ch == '_' ||
+              ch == '.')) {
+            return false;
+        }
+    }
+    return true;
+}
+
+static bool create_rendezvous_directory(const std::string& path) {
+    return mkdir(path.c_str(), 0770) == 0 || errno == EEXIST;
+}
+
+static bool nccl_sync_dir(std::string* output) {
+    if (!output) return false;
+    const char* explicit_run_id = std::getenv("RUSTRAIN_NCCL_RUN_ID");
+    const char* launch_run_id = std::getenv("RUSTRAIN_RUN_ID");
+    const char* launch_attempt_id = std::getenv("RUSTRAIN_ATTEMPT_ID");
+    const bool use_explicit_run_id = explicit_run_id && explicit_run_id[0] != '\0';
+    if (!use_explicit_run_id &&
+        (!launch_run_id || launch_run_id[0] == '\0' ||
+         !launch_attempt_id || launch_attempt_id[0] == '\0')) {
+        fprintf(stderr,
+            "[parallel_nccl] distributed initialization requires either "
+            "RUSTRAIN_NCCL_RUN_ID or both RUSTRAIN_RUN_ID and "
+            "RUSTRAIN_ATTEMPT_ID\n");
+        return false;
+    }
+
+    const char* configured_base = std::getenv("RUSTRAIN_NCCL_SYNC_DIR");
+    const char* launch_output = std::getenv("RUSTRAIN_LAUNCH_OUTPUT_DIR");
+    std::string base;
+    if (configured_base && configured_base[0] != '\0') {
+        base = configured_base;
+    } else if (launch_output && launch_output[0] != '\0') {
+        base = std::string(launch_output) + "/.rustrain-nccl";
+    } else {
+        base = "/tmp/rustrain-nccl";
+    }
+    if (!create_rendezvous_directory(base)) {
+        fprintf(stderr,
+            "[parallel_nccl] failed to create rendezvous base directory %s: %s\n",
+            base.c_str(), std::strerror(errno));
+        return false;
+    }
+
+    const char* selected_run_id =
+        use_explicit_run_id ? explicit_run_id : launch_run_id;
+    if (!valid_rendezvous_component(selected_run_id)) {
+        fprintf(stderr,
+            "[parallel_nccl] rendezvous run ID must contain only ASCII "
+            "letters, digits, '.', '_', or '-' and cannot be '.' or '..'\n");
+        return false;
+    }
+    const std::string run_id(selected_run_id);
+    std::string path = base + "/" + run_id;
+    if (!create_rendezvous_directory(path)) {
+        fprintf(stderr,
+            "[parallel_nccl] failed to create rendezvous run directory %s: %s\n",
+            path.c_str(), std::strerror(errno));
+        return false;
+    }
+    if (!use_explicit_run_id) {
+        if (!valid_rendezvous_component(launch_attempt_id)) {
+            fprintf(stderr,
+                "[parallel_nccl] rendezvous attempt ID must contain only ASCII "
+                "letters, digits, '.', '_', or '-' and cannot be '.' or '..'\n");
+            return false;
+        }
+        const std::string attempt_id(launch_attempt_id);
+        path += "/" + attempt_id;
+        if (!create_rendezvous_directory(path)) {
+            fprintf(stderr,
+                "[parallel_nccl] failed to create rendezvous attempt directory %s: %s\n",
+                path.c_str(), std::strerror(errno));
+            return false;
+        }
+    }
+    *output = std::move(path);
+    return true;
+}
+
 // ──────────────────────────────────────────────────────────────────────
 // Hand-written CUDA fused kernels (compiled from fused_kernels.cu)
 // ──────────────────────────────────────────────────────────────────────
@@ -60,6 +231,46 @@ extern "C" {
         float lr_scaled, float eps_scaled,
         float one_minus_beta1, float one_minus_beta2,
         void* stream);
+    void launch_fused_adam_multi_out_of_place(
+        void** d_src_param_ptrs, void** d_grad_ptrs,
+        float** d_src_m_ptrs, float** d_src_v_ptrs,
+        void** d_dst_param_ptrs,
+        float** d_dst_m_ptrs, float** d_dst_v_ptrs,
+        int* d_sizes, float* d_lr_scaled, float* d_eps_scaled,
+        float* d_beta1, float* d_beta2,
+        int n_params,
+        void* stream);
+    void launch_fused_multi_tensor_l2_norm(
+        void** d_grad_ptrs, int* d_sizes, int* d_groups,
+        int n_tensors, float* d_norm_squares, void* stream);
+    void launch_fused_multi_tensor_clip(
+        void** d_grad_ptrs, int* d_sizes, int* d_groups,
+        int n_tensors, const float* d_norm_squares, float max_norm,
+        void* stream);
+    void launch_fused_moe_weighted_unpermute_forward(
+        const void* returned, const void* assignment_weights,
+        const int64_t* send_index, int64_t* inverse_order, void* output,
+        int64_t assignments, int64_t tokens, int64_t hidden, int top_k,
+        int dtype, void* stream);
+    void launch_fused_moe_weighted_unpermute_backward(
+        const void* returned, const void* assignment_weights,
+        const int64_t* inverse_order, const void* grad_output,
+        void* grad_returned, void* grad_assignment_weights,
+        int64_t assignments, int64_t hidden, int top_k,
+        int dtype, void* stream);
+    void launch_fused_moe_local_permute_forward(
+        const void* received,
+        const int64_t* received_tokens, int64_t token_stride,
+        const int64_t* received_experts, int64_t expert_stride,
+        void* selected, int64_t* selected_tokens, int64_t* sorted_experts,
+        int* counts, int* offsets, int* cursors,
+        int64_t* sorted_to_received,
+        int64_t rows, int64_t hidden, int expert_count,
+        int dtype, void* stream);
+    void launch_fused_moe_local_permute_rows(
+        const void* input, const int64_t* sorted_to_received, void* output,
+        int64_t rows, int64_t hidden, int dtype, bool inverse,
+        void* stream);
 }
 
 /// Fused RMSNorm — single CUDA kernel (replaces 3 ATen ops)
@@ -83,29 +294,680 @@ static at::Tensor fused_rmsnorm_op(
 ///   d/dg = sigmoid(g) * (1 + g * (1 - sigmoid(g))) * u
 ///   d/du = silu(g) = g * sigmoid(g)
 struct NcclAllReduceFunction : public torch::autograd::Function<NcclAllReduceFunction> {
+    static ncclDataType_t dtype_for(at::ScalarType type) {
+        switch (type) {
+            case at::kBFloat16: return ncclBfloat16;
+            case at::kFloat: return ncclFloat;
+            case at::kHalf: return ncclFloat16;
+            default:
+                TORCH_CHECK(false, "unsupported NCCL all-reduce dtype: ", type);
+        }
+    }
+
+    // NCCL is normally issued on PyTorch's current stream. When an external
+    // NCCL stream is supplied by the EP runtime, fence both sides so the
+    // collective observes the producer and its result is visible to the
+    // current compute stream. This keeps the operation asynchronous without
+    // relying on a device-wide synchronize.
+    static at::Tensor allreduce(
+        const at::Tensor& input, ncclComm_t comm, cudaStream_t requested_stream,
+        ncclRedOp_t reduction = ncclSum
+    ) {
+        TORCH_CHECK(input.is_cuda(), "NCCL all-reduce requires a CUDA tensor");
+        const int dev = input.device().index();
+        cudaSetDevice(dev);
+        const auto current_stream = c10::cuda::getCurrentCUDAStream(dev).stream();
+        const auto comm_stream = requested_stream ? requested_stream : current_stream;
+        auto contiguous = input.contiguous();
+        auto output = at::empty_like(contiguous);
+
+        cudaEvent_t before = nullptr;
+        cudaEvent_t after = nullptr;
+        const bool cross_stream = comm_stream != current_stream;
+        if (cross_stream) {
+            TORCH_CHECK(cudaEventCreateWithFlags(&before, cudaEventDisableTiming) == cudaSuccess,
+                "failed to create NCCL producer event");
+            TORCH_CHECK(cudaEventCreateWithFlags(&after, cudaEventDisableTiming) == cudaSuccess,
+                "failed to create NCCL consumer event");
+            TORCH_CHECK(cudaEventRecord(before, current_stream) == cudaSuccess,
+                "failed to record NCCL producer event");
+            TORCH_CHECK(cudaStreamWaitEvent(comm_stream, before, 0) == cudaSuccess,
+                "failed to wait for NCCL producer event");
+        }
+
+        auto err = ncclAllReduce(
+            contiguous.data_ptr(), output.data_ptr(), contiguous.numel(),
+            dtype_for(contiguous.scalar_type()), reduction, comm, comm_stream);
+        TORCH_CHECK(err == ncclSuccess, "ncclAllReduce failed: ", ncclGetErrorString(err));
+
+        if (cross_stream) {
+            TORCH_CHECK(cudaEventRecord(after, comm_stream) == cudaSuccess,
+                "failed to record NCCL consumer event");
+            TORCH_CHECK(cudaStreamWaitEvent(current_stream, after, 0) == cudaSuccess,
+                "failed to wait for NCCL consumer event");
+            // Destruction is asynchronous-safe after the wait has been
+            // enqueued on the current stream.
+            cudaEventDestroy(before);
+            cudaEventDestroy(after);
+        }
+        return output;
+    }
+
     static at::Tensor forward(torch::autograd::AutogradContext* ctx,
-        at::Tensor input, int64_t comm_ptr) {
+        at::Tensor input, int64_t comm_ptr, int64_t stream_ptr) {
         ctx->saved_data["comm"] = comm_ptr;
+        ctx->saved_data["stream"] = stream_ptr;
         auto nccl_comm = reinterpret_cast<ncclComm_t>(comm_ptr);
-        int dev = input.device().index();
-        cudaSetDevice(dev);
-        auto compute_stream = c10::cuda::getCurrentCUDAStream(dev).stream();
         ctx->save_for_backward({input});
-        auto output = at::empty_like(input);
-        ncclResult_t err = ncclAllReduce(
-            input.data_ptr(), output.data_ptr(),
-            input.numel(), ncclBfloat16, ncclSum,
-            nccl_comm, compute_stream);
-        if (err != ncclSuccess) {
-            fprintf(stderr, "[ep] ncclAllReduce fwd FAILED: %d (%s) dev=%d\n", err, ncclGetErrorString(err), dev);
-        }
-        return output;
+        return allreduce(input, nccl_comm,
+            reinterpret_cast<cudaStream_t>(stream_ptr));
     }
     static std::vector<at::Tensor> backward(torch::autograd::AutogradContext* ctx,
         std::vector<at::Tensor> grad_output) {
         // Expert weights frozen — gradient flows through residual connection.
         // Save input in forward so autograd keeps it alive until backward.
-        return {grad_output[0], at::Tensor()};
+        return {grad_output[0], at::Tensor(), at::Tensor()};
+    }
+};
+
+// Move an NCCL payload to the dedicated communication stream without a device
+// synchronize. The producer event is recorded on the compute stream before
+// NCCL launch; the consumer event makes received data visible to compute.
+// When no external stream is configured this is a zero-cost no-op.
+struct Qwen36StreamFence {
+    cudaStream_t current;
+    cudaStream_t operation;
+    cudaEvent_t before = nullptr;
+    cudaEvent_t after = nullptr;
+    bool cross_stream = false;
+
+    Qwen36StreamFence(cudaStream_t current_stream, cudaStream_t requested_stream)
+        : current(current_stream),
+          operation(requested_stream ? requested_stream : current_stream),
+          cross_stream(operation != current) {
+        if (!cross_stream) return;
+        TORCH_CHECK(cudaEventCreateWithFlags(
+                &before, cudaEventDisableTiming) == cudaSuccess,
+            "failed to create NCCL producer event");
+        TORCH_CHECK(cudaEventCreateWithFlags(
+                &after, cudaEventDisableTiming) == cudaSuccess,
+            "failed to create NCCL consumer event");
+        TORCH_CHECK(cudaEventRecord(before, current) == cudaSuccess,
+            "failed to record NCCL producer event");
+        TORCH_CHECK(cudaStreamWaitEvent(operation, before, 0) == cudaSuccess,
+            "failed to fence NCCL operation stream");
+    }
+
+    void complete() {
+        if (!cross_stream) return;
+        TORCH_CHECK(cudaEventRecord(after, operation) == cudaSuccess,
+            "failed to record NCCL consumer event");
+        TORCH_CHECK(cudaStreamWaitEvent(current, after, 0) == cudaSuccess,
+            "failed to fence compute stream after NCCL");
+        cudaEventDestroy(before);
+        cudaEventDestroy(after);
+        before = nullptr;
+        after = nullptr;
+        cross_stream = false;
+    }
+
+    ~Qwen36StreamFence() {
+        if (before) cudaEventDestroy(before);
+        if (after) cudaEventDestroy(after);
+    }
+};
+
+// Lifetime guard for the optional MoE shared-expert overlap path. CUDA event
+// destruction is safe while work is pending; the explicit consumer wait is
+// still enqueued before the result is consumed on the compute stream.
+struct Qwen36EventPair {
+    cudaEvent_t ready = nullptr;
+    cudaEvent_t done = nullptr;
+
+    ~Qwen36EventPair() {
+        if (ready) cudaEventDestroy(ready);
+        if (done) cudaEventDestroy(done);
+    }
+};
+
+// Megatron's copy_to_tensor_model_parallel_region equivalent: replicated
+// input in forward, sum the column-parallel input-gradient contributions in
+// backward before they flow into the preceding replicated sub-layer.
+struct TpCopyToRegionFunction : public torch::autograd::Function<TpCopyToRegionFunction> {
+    static at::Tensor forward(torch::autograd::AutogradContext* ctx,
+        at::Tensor input, int64_t comm_ptr, int64_t stream_ptr) {
+        ctx->saved_data["comm"] = comm_ptr;
+        ctx->saved_data["stream"] = stream_ptr;
+        return input;
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output) {
+        auto comm = reinterpret_cast<ncclComm_t>(ctx->saved_data["comm"].toInt());
+        auto stream = reinterpret_cast<cudaStream_t>(ctx->saved_data["stream"].toInt());
+        auto grad_input = NcclAllReduceFunction::allreduce(
+            grad_output[0], comm, stream);
+        return {grad_input, at::Tensor(), at::Tensor()};
+    }
+};
+
+// Sequence-parallel collectives keep the public activation layout
+// [batch, sequence, hidden]. NCCL's all-gather/reduce-scatter layout is
+// rank-major, so batch>1 requires an explicit permutation at both boundaries.
+static at::Tensor tp_sequence_all_gather(
+    const at::Tensor& input, ncclComm_t comm, cudaStream_t requested_stream,
+    int64_t world
+) {
+    TORCH_CHECK(input.is_cuda() && input.dim() == 3 && input.is_contiguous(),
+        "sequence all-gather requires contiguous CUDA [batch, seq, hidden]");
+    TORCH_CHECK(comm && world > 1, "sequence all-gather requires TP communicator");
+    const int device = input.device().index();
+    const auto current = c10::cuda::getCurrentCUDAStream(device).stream();
+    Qwen36StreamFence fence(current, requested_stream);
+    auto rank_major = at::empty(
+        {world, input.size(0), input.size(1), input.size(2)}, input.options());
+    const auto error = ncclAllGather(
+        input.data_ptr(), rank_major.data_ptr(), input.numel(),
+        NcclAllReduceFunction::dtype_for(input.scalar_type()), comm,
+        fence.operation);
+    TORCH_CHECK(error == ncclSuccess, "TP sequence all-gather failed: ",
+        ncclGetErrorString(error));
+    fence.complete();
+    return rank_major.permute({1, 0, 2, 3}).reshape({
+        input.size(0), input.size(1) * world, input.size(2)}).contiguous();
+}
+
+static at::Tensor tp_sequence_reduce_scatter(
+    const at::Tensor& input, ncclComm_t comm, cudaStream_t requested_stream,
+    int64_t world
+) {
+    TORCH_CHECK(input.is_cuda() && input.dim() == 3,
+        "sequence reduce-scatter requires CUDA [batch, seq, hidden]");
+    TORCH_CHECK(comm && world > 1 && input.size(1) % world == 0,
+        "sequence reduce-scatter requires sequence divisible by TP_SIZE");
+    const int64_t local_sequence = input.size(1) / world;
+    auto rank_major = input.reshape({
+        input.size(0), world, local_sequence, input.size(2)})
+        .permute({1, 0, 2, 3}).contiguous();
+    auto output = at::empty(
+        {input.size(0), local_sequence, input.size(2)}, input.options());
+    const int device = input.device().index();
+    const auto current = c10::cuda::getCurrentCUDAStream(device).stream();
+    Qwen36StreamFence fence(current, requested_stream);
+    const auto error = ncclReduceScatter(
+        rank_major.data_ptr(), output.data_ptr(), output.numel(),
+        NcclAllReduceFunction::dtype_for(input.scalar_type()), ncclSum, comm,
+        fence.operation);
+    TORCH_CHECK(error == ncclSuccess, "TP sequence reduce-scatter failed: ",
+        ncclGetErrorString(error));
+    fence.complete();
+    return output;
+}
+
+// Ring attention communication primitive.  Every rank issues the same
+// send/recv schedule, including the final rotation which returns the local
+// KV block to its owner.  Keeping this primitive separate from attention
+// makes it possible to reuse the exact schedule in the hand-written backward.
+static std::vector<at::Tensor> cp_ring_exchange(
+    const std::vector<at::Tensor>& payloads, ncclComm_t comm,
+    cudaStream_t requested_stream, int64_t rank, int64_t world
+) {
+    TORCH_CHECK(!payloads.empty() && comm && world > 1 &&
+            rank >= 0 && rank < world,
+        "CP ring exchange requires a valid communicator and world");
+    const auto& first = payloads.front();
+    TORCH_CHECK(first.is_cuda() && first.is_contiguous(),
+        "CP ring payloads must be contiguous CUDA tensors");
+    const int device = first.device().index();
+    const auto current = c10::cuda::getCurrentCUDAStream(device).stream();
+    Qwen36StreamFence fence(current, requested_stream);
+    std::vector<at::Tensor> received;
+    received.reserve(payloads.size());
+    for (const auto& payload : payloads) {
+        TORCH_CHECK(payload.is_cuda() && payload.is_contiguous() &&
+                payload.sizes() == first.sizes(),
+            "CP ring payloads must have identical contiguous CUDA shapes");
+        received.push_back(at::empty_like(payload));
+    }
+    const int peer_next = static_cast<int>((rank + 1) % world);
+    const int peer_prev = static_cast<int>((rank + world - 1) % world);
+    TORCH_CHECK(ncclGroupStart() == ncclSuccess,
+        "CP ring group start failed");
+    ncclResult_t first_error = ncclSuccess;
+    for (const auto& payload : payloads) {
+        const auto error = ncclSend(
+            payload.data_ptr(), payload.numel(),
+            NcclAllReduceFunction::dtype_for(payload.scalar_type()),
+            peer_next, comm, fence.operation);
+        if (first_error == ncclSuccess && error != ncclSuccess)
+            first_error = error;
+    }
+    for (auto& output : received) {
+        const auto error = ncclRecv(
+            output.data_ptr(), output.numel(),
+            NcclAllReduceFunction::dtype_for(output.scalar_type()),
+            peer_prev, comm, fence.operation);
+        if (first_error == ncclSuccess && error != ncclSuccess)
+            first_error = error;
+    }
+    const auto group_error = ncclGroupEnd();
+    TORCH_CHECK(first_error == ncclSuccess,
+        "CP ring P2P failed: ", ncclGetErrorString(first_error));
+    TORCH_CHECK(group_error == ncclSuccess,
+        "CP ring group end failed: ", ncclGetErrorString(group_error));
+    fence.complete();
+    return received;
+}
+
+static at::Tensor cp_ring_attention_mask(
+    const at::Tensor& attention_mask, const at::Tensor& device_tensor,
+    int64_t batch, int64_t query_start,
+    int64_t query_length, int64_t key_start, int64_t key_length
+) {
+    auto options = at::TensorOptions().dtype(at::kLong)
+        .device(device_tensor.device());
+    auto query_positions = at::arange(
+        query_start, query_start + query_length, options).view({1, 1, query_length, 1});
+    auto key_positions = at::arange(
+        key_start, key_start + key_length, options).view({1, 1, 1, key_length});
+    auto mask = key_positions.gt(query_positions);
+    if (attention_mask.defined() && attention_mask.numel() > 0) {
+        TORCH_CHECK(attention_mask.dim() == 2 &&
+                attention_mask.size(0) == batch,
+            "CP ring attention mask must be [batch, global_seq]");
+        auto key_valid = attention_mask.to(at::kBool).narrow(
+            1, key_start, key_length).view({batch, 1, 1, key_length});
+        mask = mask.logical_or(key_valid.logical_not());
+    }
+    return mask.expand({batch, 1, query_length, key_length});
+}
+
+// Memory-scalable causal full attention for arbitrary CP world sizes.  The
+// forward maintains online-softmax state while KV blocks rotate through the
+// ring.  Backward replays the same ring and carries FP32 dK/dV accumulators
+// with each block so every owner receives contributions from all query ranks.
+struct Qwen36RingAttentionFunction
+    : public torch::autograd::Function<Qwen36RingAttentionFunction> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx, at::Tensor q,
+        at::Tensor local_kv, at::Tensor attention_mask, int64_t comm_ptr,
+        int64_t stream_ptr, int64_t rank, int64_t world, int64_t group,
+        int64_t query_start
+    ) {
+        TORCH_CHECK(q.dim() == 4 && local_kv.dim() == 4 &&
+                q.is_cuda() && local_kv.is_cuda() && q.is_contiguous() &&
+                local_kv.is_contiguous(),
+            "CP ring attention requires contiguous CUDA Q/KV tensors");
+        TORCH_CHECK(world > 1 && group > 0 && q.size(1) == local_kv.size(1) * group &&
+                q.size(2) == local_kv.size(2),
+            "CP ring attention Q/KV head geometry is invalid");
+        TORCH_CHECK(local_kv.size(3) % 2 == 0,
+            "CP ring KV payload must concatenate K and V");
+        const int64_t batch = q.size(0);
+        const int64_t query_length = q.size(2);
+        const int64_t head_dim = local_kv.size(3) / 2;
+        const double scale = 1.0 / std::sqrt(static_cast<double>(head_dim));
+        auto q_float = q.to(at::kFloat);
+        auto running_m = at::full(
+            {batch, q.size(1), query_length, 1}, -std::numeric_limits<float>::infinity(),
+            q_float.options());
+        auto running_l = at::zeros_like(running_m);
+        auto running_o = at::zeros(
+            {batch, q.size(1), query_length, head_dim}, q_float.options());
+        auto current_kv = local_kv;
+        for (int64_t step = 0; step < world; ++step) {
+            const int64_t owner = (rank + world - step) % world;
+            auto k = current_kv.narrow(-1, 0, head_dim)
+                .to(at::kFloat).repeat_interleave(group, 1);
+            auto v = current_kv.narrow(-1, head_dim, head_dim)
+                .to(at::kFloat).repeat_interleave(group, 1);
+            auto scores = at::matmul(q_float, k.transpose(-2, -1)) * scale;
+            auto mask = cp_ring_attention_mask(
+                attention_mask, q, batch, query_start, query_length,
+                owner * query_length, query_length);
+            scores = scores.masked_fill(mask, -std::numeric_limits<float>::infinity());
+            auto block_m = std::get<0>(scores.max(-1, true));
+            auto new_m = at::maximum(running_m, block_m);
+            auto new_m_safe = at::where(
+                at::isfinite(new_m), new_m, at::zeros_like(new_m));
+            auto old_scale = at::where(
+                at::isfinite(running_m), at::exp(running_m - new_m_safe),
+                at::zeros_like(running_m));
+            auto block_exp = at::exp(scores - new_m_safe).masked_fill(mask, 0.0);
+            running_o = running_o * old_scale + at::matmul(block_exp, v);
+            running_l = running_l * old_scale + block_exp.sum(-1, true);
+            running_m = new_m;
+            auto rotated = cp_ring_exchange(
+                {current_kv}, reinterpret_cast<ncclComm_t>(comm_ptr),
+                reinterpret_cast<cudaStream_t>(stream_ptr), rank, world);
+            current_kv = rotated.front();
+        }
+        auto output = running_o / running_l.clamp_min(1e-9);
+        auto lse = at::where(
+            running_l > 0, running_m + at::log(running_l),
+            at::full_like(running_l, -std::numeric_limits<float>::infinity()));
+        ctx->saved_data["comm"] = comm_ptr;
+        ctx->saved_data["stream"] = stream_ptr;
+        ctx->saved_data["rank"] = rank;
+        ctx->saved_data["world"] = world;
+        ctx->saved_data["group"] = group;
+        ctx->saved_data["query_start"] = query_start;
+        ctx->save_for_backward({q, local_kv, attention_mask, output, lse});
+        return output.to(q.scalar_type());
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output
+    ) {
+        auto saved = ctx->get_saved_variables();
+        auto q = saved[0];
+        auto local_kv = saved[1];
+        auto attention_mask = saved[2];
+        auto output = saved[3];
+        auto lse = saved[4];
+        const int64_t rank = ctx->saved_data["rank"].toInt();
+        const int64_t world = ctx->saved_data["world"].toInt();
+        const int64_t group = ctx->saved_data["group"].toInt();
+        const int64_t query_start = ctx->saved_data["query_start"].toInt();
+        const int64_t head_dim = local_kv.size(3) / 2;
+        const int64_t batch = q.size(0);
+        const int64_t query_length = q.size(2);
+        const double scale = 1.0 / std::sqrt(static_cast<double>(head_dim));
+        auto q_float = q.to(at::kFloat);
+        auto grad_out = grad_output[0].to(at::kFloat);
+        auto grad_q = at::zeros_like(q_float);
+        auto grad_kv = at::zeros(local_kv.sizes(), q_float.options());
+        auto current_kv = local_kv;
+        auto current_grad = at::zeros_like(grad_kv);
+        auto delta = (grad_out * output).sum(-1, true);
+        for (int64_t step = 0; step < world; ++step) {
+            const int64_t owner = (rank + world - step) % world;
+            auto k = current_kv.narrow(-1, 0, head_dim)
+                .to(at::kFloat).repeat_interleave(group, 1);
+            auto v = current_kv.narrow(-1, head_dim, head_dim)
+                .to(at::kFloat).repeat_interleave(group, 1);
+            auto scores = at::matmul(q_float, k.transpose(-2, -1)) * scale;
+            auto mask = cp_ring_attention_mask(
+                attention_mask, q, batch, query_start, query_length,
+                owner * query_length, query_length);
+            auto safe_lse = at::where(
+                at::isfinite(lse), lse, at::zeros_like(lse));
+            scores = scores.masked_fill(mask, -std::numeric_limits<float>::infinity());
+            auto probs = at::exp(scores - safe_lse);
+            probs = probs.masked_fill(mask, 0.0);
+            auto grad_scores = probs * (
+                at::matmul(grad_out, v.transpose(-2, -1)) - delta);
+            grad_q = grad_q + at::matmul(grad_scores, k) * scale;
+            auto grad_k_expanded = at::matmul(
+                grad_scores.transpose(-2, -1), q_float) * scale;
+            auto grad_v_expanded = at::matmul(
+                probs.transpose(-2, -1), grad_out);
+            auto grad_k = grad_k_expanded.view({
+                batch, local_kv.size(1), group, query_length, head_dim})
+                .sum(2);
+            auto grad_v = grad_v_expanded.view({
+                batch, local_kv.size(1), group, query_length, head_dim})
+                .sum(2);
+            current_grad.narrow(-1, 0, head_dim).add_(grad_k);
+            current_grad.narrow(-1, head_dim, head_dim).add_(grad_v);
+            auto rotated = cp_ring_exchange(
+                {current_kv, current_grad},
+                reinterpret_cast<ncclComm_t>(ctx->saved_data["comm"].toInt()),
+                reinterpret_cast<cudaStream_t>(ctx->saved_data["stream"].toInt()),
+                rank, world);
+            current_kv = rotated[0];
+            current_grad = rotated[1];
+        }
+        grad_kv.copy_(current_grad);
+        return {grad_q.to(q.scalar_type()),
+            grad_kv.to(local_kv.scalar_type()), at::Tensor(), at::Tensor(),
+            at::Tensor(), at::Tensor(), at::Tensor(), at::Tensor(), at::Tensor(),
+        };
+    }
+};
+
+// Megatron gather_from_sequence_parallel_region: gather in forward and sum
+// column-parallel dgrad contributions while scattering them in backward.
+struct TpGatherFromSequenceFunction
+    : public torch::autograd::Function<TpGatherFromSequenceFunction> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx, at::Tensor input,
+        int64_t comm_ptr, int64_t stream_ptr, int64_t world
+    ) {
+        ctx->saved_data["comm"] = comm_ptr;
+        ctx->saved_data["stream"] = stream_ptr;
+        ctx->saved_data["world"] = world;
+        return tp_sequence_all_gather(
+            input.contiguous(), reinterpret_cast<ncclComm_t>(comm_ptr),
+            reinterpret_cast<cudaStream_t>(stream_ptr), world);
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output
+    ) {
+        auto grad_input = tp_sequence_reduce_scatter(
+            grad_output[0],
+            reinterpret_cast<ncclComm_t>(ctx->saved_data["comm"].toInt()),
+            reinterpret_cast<cudaStream_t>(ctx->saved_data["stream"].toInt()),
+            ctx->saved_data["world"].toInt());
+        return {grad_input, at::Tensor(), at::Tensor(), at::Tensor()};
+    }
+};
+
+// Megatron reduce_scatter_to_sequence_parallel_region: sum row-parallel
+// outputs in forward; backward is a plain all-gather because each rank owns a
+// disjoint upstream sequence slice.
+struct TpReduceScatterToSequenceFunction
+    : public torch::autograd::Function<TpReduceScatterToSequenceFunction> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx, at::Tensor input,
+        int64_t comm_ptr, int64_t stream_ptr, int64_t world
+    ) {
+        ctx->saved_data["comm"] = comm_ptr;
+        ctx->saved_data["stream"] = stream_ptr;
+        ctx->saved_data["world"] = world;
+        return tp_sequence_reduce_scatter(
+            input, reinterpret_cast<ncclComm_t>(comm_ptr),
+            reinterpret_cast<cudaStream_t>(stream_ptr), world);
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output
+    ) {
+        auto grad_input = tp_sequence_all_gather(
+            grad_output[0].contiguous(),
+            reinterpret_cast<ncclComm_t>(ctx->saved_data["comm"].toInt()),
+            reinterpret_cast<cudaStream_t>(ctx->saved_data["stream"].toInt()),
+            ctx->saved_data["world"].toInt());
+        return {grad_input, at::Tensor(), at::Tensor(), at::Tensor()};
+    }
+};
+
+// The vocabulary loss already sums its hidden dgrad over TP. Its gather must
+// therefore map backward to a plain local sequence slice, not another
+// reduce-scatter, which would multiply the hidden gradient by TP_SIZE.
+struct TpGatherForLossFunction
+    : public torch::autograd::Function<TpGatherForLossFunction> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx, at::Tensor input,
+        int64_t comm_ptr, int64_t stream_ptr, int64_t rank, int64_t world
+    ) {
+        ctx->saved_data["rank"] = rank;
+        ctx->saved_data["world"] = world;
+        ctx->saved_data["local_sequence"] = input.size(1);
+        return tp_sequence_all_gather(
+            input.contiguous(), reinterpret_cast<ncclComm_t>(comm_ptr),
+            reinterpret_cast<cudaStream_t>(stream_ptr), world);
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output
+    ) {
+        const int64_t rank = ctx->saved_data["rank"].toInt();
+        const int64_t local_sequence =
+            ctx->saved_data["local_sequence"].toInt();
+        auto grad_input = grad_output[0]
+            .narrow(1, rank * local_sequence, local_sequence).contiguous();
+        return {grad_input, at::Tensor(), at::Tensor(), at::Tensor(), at::Tensor()};
+    }
+};
+
+// Headwise GDN context parallelism maps local sequence slices to local head
+// bundles before the recurrent kernel, then performs the inverse mapping
+// before the output projection. Inputs are packed peer-major on the feature
+// axis so section-aware Q/K/V slicing stays outside the communication layer.
+static at::Tensor cp_sequence_to_head_exchange(
+    const at::Tensor& input, ncclComm_t comm,
+    cudaStream_t requested_stream, int64_t world
+) {
+    TORCH_CHECK(input.is_cuda() && input.dim() == 3 && input.is_contiguous(),
+        "CP sequence-to-head requires contiguous CUDA [batch, seq, features]");
+    TORCH_CHECK(comm && world == 2 && input.size(2) % world == 0,
+        "CP sequence-to-head currently requires CP_SIZE=2 and divisible features");
+    const int64_t peer_features = input.size(2) / world;
+    auto packed = input.reshape({
+        input.size(0), input.size(1), world, peer_features})
+        .permute({2, 0, 1, 3}).contiguous();
+    auto received = at::empty_like(packed);
+    const int device = input.device().index();
+    const auto current = c10::cuda::getCurrentCUDAStream(device).stream();
+    Qwen36StreamFence fence(current, requested_stream);
+    TORCH_CHECK(ncclGroupStart() == ncclSuccess,
+        "CP sequence-to-head group start failed");
+    ncclResult_t first_error = ncclSuccess;
+    for (int peer = 0; peer < world; ++peer) {
+        auto source = packed.select(0, peer);
+        const auto error = ncclSend(
+            source.data_ptr(), source.numel(),
+            NcclAllReduceFunction::dtype_for(input.scalar_type()), peer,
+            comm, fence.operation);
+        if (first_error == ncclSuccess && error != ncclSuccess)
+            first_error = error;
+    }
+    for (int peer = 0; peer < world; ++peer) {
+        auto destination = received.select(0, peer);
+        const auto error = ncclRecv(
+            destination.data_ptr(), destination.numel(),
+            NcclAllReduceFunction::dtype_for(input.scalar_type()), peer,
+            comm, fence.operation);
+        if (first_error == ncclSuccess && error != ncclSuccess)
+            first_error = error;
+    }
+    const auto group_error = ncclGroupEnd();
+    TORCH_CHECK(first_error == ncclSuccess,
+        "CP sequence-to-head P2P failed: ", ncclGetErrorString(first_error));
+    TORCH_CHECK(group_error == ncclSuccess,
+        "CP sequence-to-head group end failed: ",
+        ncclGetErrorString(group_error));
+    fence.complete();
+    return received.permute({1, 0, 2, 3}).reshape({
+        input.size(0), input.size(1) * world, peer_features}).contiguous();
+}
+
+static at::Tensor cp_head_to_sequence_exchange(
+    const at::Tensor& input, ncclComm_t comm,
+    cudaStream_t requested_stream, int64_t world
+) {
+    TORCH_CHECK(input.is_cuda() && input.dim() == 3 && input.is_contiguous(),
+        "CP head-to-sequence requires contiguous CUDA [batch, seq, features]");
+    TORCH_CHECK(comm && world == 2 && input.size(1) % world == 0,
+        "CP head-to-sequence currently requires CP_SIZE=2 and divisible sequence");
+    const int64_t local_sequence = input.size(1) / world;
+    auto packed = input.reshape({
+        input.size(0), world, local_sequence, input.size(2)})
+        .permute({1, 0, 2, 3}).contiguous();
+    auto received = at::empty_like(packed);
+    const int device = input.device().index();
+    const auto current = c10::cuda::getCurrentCUDAStream(device).stream();
+    Qwen36StreamFence fence(current, requested_stream);
+    TORCH_CHECK(ncclGroupStart() == ncclSuccess,
+        "CP head-to-sequence group start failed");
+    ncclResult_t first_error = ncclSuccess;
+    for (int peer = 0; peer < world; ++peer) {
+        auto source = packed.select(0, peer);
+        const auto error = ncclSend(
+            source.data_ptr(), source.numel(),
+            NcclAllReduceFunction::dtype_for(input.scalar_type()), peer,
+            comm, fence.operation);
+        if (first_error == ncclSuccess && error != ncclSuccess)
+            first_error = error;
+    }
+    for (int peer = 0; peer < world; ++peer) {
+        auto destination = received.select(0, peer);
+        const auto error = ncclRecv(
+            destination.data_ptr(), destination.numel(),
+            NcclAllReduceFunction::dtype_for(input.scalar_type()), peer,
+            comm, fence.operation);
+        if (first_error == ncclSuccess && error != ncclSuccess)
+            first_error = error;
+    }
+    const auto group_error = ncclGroupEnd();
+    TORCH_CHECK(first_error == ncclSuccess,
+        "CP head-to-sequence P2P failed: ", ncclGetErrorString(first_error));
+    TORCH_CHECK(group_error == ncclSuccess,
+        "CP head-to-sequence group end failed: ",
+        ncclGetErrorString(group_error));
+    fence.complete();
+    return received.permute({1, 2, 0, 3}).reshape({
+        input.size(0), local_sequence, input.size(2) * world}).contiguous();
+}
+
+struct CpSequenceToHeadFunction
+    : public torch::autograd::Function<CpSequenceToHeadFunction> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx, at::Tensor input,
+        int64_t comm_ptr, int64_t stream_ptr, int64_t world
+    ) {
+        ctx->saved_data["comm"] = comm_ptr;
+        ctx->saved_data["stream"] = stream_ptr;
+        ctx->saved_data["world"] = world;
+        return cp_sequence_to_head_exchange(
+            input.contiguous(), reinterpret_cast<ncclComm_t>(comm_ptr),
+            reinterpret_cast<cudaStream_t>(stream_ptr), world);
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output
+    ) {
+        auto grad_input = cp_head_to_sequence_exchange(
+            grad_output[0].contiguous(),
+            reinterpret_cast<ncclComm_t>(ctx->saved_data["comm"].toInt()),
+            reinterpret_cast<cudaStream_t>(ctx->saved_data["stream"].toInt()),
+            ctx->saved_data["world"].toInt());
+        return {grad_input, at::Tensor(), at::Tensor(), at::Tensor()};
+    }
+};
+
+struct CpHeadToSequenceFunction
+    : public torch::autograd::Function<CpHeadToSequenceFunction> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx, at::Tensor input,
+        int64_t comm_ptr, int64_t stream_ptr, int64_t world
+    ) {
+        ctx->saved_data["comm"] = comm_ptr;
+        ctx->saved_data["stream"] = stream_ptr;
+        ctx->saved_data["world"] = world;
+        return cp_head_to_sequence_exchange(
+            input.contiguous(), reinterpret_cast<ncclComm_t>(comm_ptr),
+            reinterpret_cast<cudaStream_t>(stream_ptr), world);
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output
+    ) {
+        auto grad_input = cp_sequence_to_head_exchange(
+            grad_output[0].contiguous(),
+            reinterpret_cast<ncclComm_t>(ctx->saved_data["comm"].toInt()),
+            reinterpret_cast<cudaStream_t>(ctx->saved_data["stream"].toInt()),
+            ctx->saved_data["world"].toInt());
+        return {grad_input, at::Tensor(), at::Tensor(), at::Tensor()};
     }
 };
 
@@ -139,7 +1001,8 @@ struct FusedSwiGLUFunction : public torch::autograd::Function<FusedSwiGLUFunctio
 static at::Tensor fused_swiglu_op(
     const at::Tensor& gate_out, const at::Tensor& up_out, double limit
 ) {
-    // ATen ops (keeps autograd graph) — TODO: re-enable CUDA kernel with autograd
+    // Keep the ATen path here because its backward is already part of the
+    // autograd graph; the surrounding matmuls remain fused by the layer path.
     auto inter = at::silu(gate_out) * up_out;
     if (limit > 0.0) inter = inter.clamp(-limit, limit);
     return inter;
@@ -334,13 +1197,16 @@ static at::Tensor full_attention(
         auto attn_out = at::scaled_dot_product_attention(
             q, k, v, additive_mask, 0.0, false, c10::nullopt, true  // is_causal=false, enable_gqa=true
         );
-        return attn_out.transpose(1, 2).reshape({batch, seq, qkv_dim}).matmul(o_proj.t());
+        // Qwen3.5/3.6 full attention gates the attention value before o_proj.
+        // Applying the gate after o_proj is not equivalent when o_proj mixes features.
+        auto gated_attn = attn_out * at::sigmoid(gate).to(attn_out.scalar_type());
+        return gated_attn.transpose(1, 2).reshape({batch, seq, qkv_dim}).matmul(o_proj.t());
     } else {
         auto attn_out = at::scaled_dot_product_attention(
             q, k, v, c10::nullopt, 0.0, true, c10::nullopt, true  // is_causal=true, enable_gqa=true
         );
-        auto result = attn_out.transpose(1, 2).reshape({batch, seq, qkv_dim}).matmul(o_proj.t());
-        result = result * at::sigmoid(gate).to(result.scalar_type());
+        auto gated_attn = attn_out * at::sigmoid(gate).to(attn_out.scalar_type());
+        auto result = gated_attn.transpose(1, 2).reshape({batch, seq, qkv_dim}).matmul(o_proj.t());
         gate = at::Tensor();
         return result;
     }
@@ -351,23 +1217,274 @@ static at::Tensor full_attention(
 // ──────────────────────────────────────────────────────────────────────
 
 // Forward declaration for CUDA kernel (defined in delta_rule.cu)
-extern "C" void cuda_gated_delta_rule(
+extern "C" int cuda_gated_delta_rule(
     const float* q, const float* k, const float* v,
     const float* g_exp, const float* beta,
     float* state, float* out, float* delta_buf,
-    int BH, int seq_len, int key_dim, int val_dim
+    float* state_checkpoints, int checkpoint_stride,
+    const int32_t* lengths, int heads_per_batch,
+    int BH, int seq_len, int key_dim, int val_dim, void* stream
 );
 
-extern "C" void cuda_gated_delta_rule_backward(
+extern "C" int cuda_gated_delta_rule_backward(
     const float* q, const float* k, const float* v,
     const float* g_exp, const float* beta,
     const float* final_state, const float* delta_buf,
+    const float* state_checkpoints, int checkpoint_stride,
     const float* grad_out,
     float* grad_q, float* grad_k, float* grad_v,
     float* grad_g, float* grad_beta,
-    int BH, int seq_len, int key_dim, int val_dim
+    const int32_t* lengths, int heads_per_batch,
+    int BH, int seq_len, int key_dim, int val_dim, void* stream
 );
 
+// Correctness reference for the gated delta rule. This deliberately uses
+// ATen batched operations inside C++, so it remains outside the Rust hot path
+// while providing a complete autograd oracle for the custom CUDA forward.
+static at::Tensor gated_delta_rule_reference(
+    const at::Tensor& q, const at::Tensor& k, const at::Tensor& v,
+    const at::Tensor& g_exp, const at::Tensor& beta,
+    const at::Tensor& lengths, int64_t heads_per_batch
+) {
+    TORCH_CHECK(q.dim() == 3 && k.dim() == 3 && v.dim() == 3,
+        "gated_delta_rule_reference expects [BH, S, D] tensors");
+    const int64_t bh = q.size(0);
+    const int64_t seq = q.size(1);
+    const int64_t key_dim = q.size(2);
+    const int64_t val_dim = v.size(2);
+    TORCH_CHECK(k.size(0) == bh && k.size(1) == seq && k.size(2) == key_dim,
+        "q/k shape mismatch");
+    TORCH_CHECK(v.size(0) == bh && v.size(1) == seq,
+        "v shape mismatch");
+    TORCH_CHECK(g_exp.size(0) == bh && g_exp.size(1) == seq &&
+                beta.size(0) == bh && beta.size(1) == seq,
+        "g/beta shape mismatch");
+    if (lengths.defined()) {
+        TORCH_CHECK(lengths.dim() == 1 && lengths.scalar_type() == at::kInt &&
+                    lengths.size(0) * heads_per_batch == bh,
+            "lengths/head mapping mismatch");
+    }
+
+    auto state = at::zeros({bh, key_dim, val_dim},
+        q.options().dtype(at::kFloat));
+    auto qf = q.to(at::kFloat);
+    auto kf = k.to(at::kFloat);
+    auto vf = v.to(at::kFloat);
+    auto gf = g_exp.to(at::kFloat);
+    auto bf = beta.to(at::kFloat);
+    auto bh_lengths = lengths.defined()
+        ? lengths.repeat_interleave(heads_per_batch)
+        : at::Tensor();
+    std::vector<at::Tensor> outputs;
+    outputs.reserve(seq);
+    for (int64_t t = 0; t < seq; ++t) {
+        auto gt = gf.select(1, t).view({bh, 1, 1});
+        auto kt = kf.select(1, t);
+        auto vt = vf.select(1, t);
+        auto qt = qf.select(1, t);
+        auto bt = bf.select(1, t).view({bh, 1});
+        auto active = bh_lengths.defined()
+            ? (bh_lengths > t).to(at::kFloat).view({bh, 1, 1})
+            : at::Tensor();
+        state = active.defined()
+            ? state * (gt * active + (1.0 - active))
+            : state * gt;
+        auto kv = at::bmm(kt.unsqueeze(1), state).squeeze(1);
+        auto delta = (vt - kv) * bt;
+        if (active.defined()) delta = delta * active.view({bh, 1});
+        state = state + kt.unsqueeze(2) * delta.unsqueeze(1);
+        auto output = at::bmm(qt.unsqueeze(1), state).squeeze(1);
+        if (active.defined()) output = output * active.view({bh, 1});
+        outputs.push_back(output);
+    }
+    return at::stack(outputs, 1);
+}
+
+// Autograd requires every Function input to be a defined tensor with a
+// device. An empty device-local int32 tensor is the dense-sequence sentinel;
+// CUDA launchers still receive a null lengths pointer for this case.
+static at::Tensor gdn_lengths_arg(
+    const at::Tensor& lengths, const at::Tensor& reference
+) {
+    return lengths.defined()
+        ? lengths
+        : at::empty({0}, reference.options().dtype(at::kInt));
+}
+
+static bool try_parse_gdn_state_checkpoint_config(int64_t* parsed_out) {
+    if (!parsed_out) return false;
+    const char* value = std::getenv("QWEN36_GDN_STATE_CHECKPOINT_STRIDE");
+    if (!value || value[0] == '\0') {
+        *parsed_out = -1;
+        return true;
+    }
+    errno = 0;
+    char* end = nullptr;
+    const long long parsed = std::strtoll(value, &end, 10);
+    if (errno == ERANGE || !end || *end != '\0' || parsed < 0 ||
+            parsed == 1)
+        return false;
+    *parsed_out = static_cast<int64_t>(parsed);
+    return true;
+}
+
+static bool gdn_state_checkpoint_environment_valid() {
+    int64_t parsed = 0;
+    if (!try_parse_gdn_state_checkpoint_config(&parsed)) return false;
+    if (parsed != 0) return true;
+    // Zero disables checkpoint storage and is only valid for diagnostic
+    // inverse/reference backward, never for production replay/chunkwise paths.
+    return !env_enabled("QWEN36_GDN_CHUNKWISE_BWD") &&
+        (env_enabled("QWEN36_GDN_INVERSE_BWD") ||
+         env_enabled("QWEN36_DELTA_REFERENCE_BWD"));
+}
+
+static int64_t gdn_state_checkpoint_config() {
+    int64_t parsed = 0;
+    TORCH_CHECK(try_parse_gdn_state_checkpoint_config(&parsed),
+        "QWEN36_GDN_STATE_CHECKPOINT_STRIDE must be unset, zero, or an "
+        "integer of at least 2");
+    return parsed;
+}
+
+static int64_t gdn_state_checkpoint_stride(int64_t seq) {
+    const int64_t configured = gdn_state_checkpoint_config();
+    if (configured < 0) {
+        if (seq <= 1) return 0;
+        const auto automatic = static_cast<int64_t>(
+            std::ceil(std::sqrt(static_cast<double>(seq))));
+        return std::max<int64_t>(2, std::min<int64_t>(automatic, seq - 1));
+    }
+    if (configured == 0 || seq <= 0) return 0;
+    return std::min(configured, seq);
+}
+
+// The forward and backward CUDA kernels are wrapped in one autograd Function.
+// Set QWEN36_DELTA_REFERENCE_BWD=1 to run the ATen recurrence oracle for parity
+// debugging; production training uses the current-stream fused backward.
+struct GatedDeltaRuleFunction : public torch::autograd::Function<GatedDeltaRuleFunction> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx,
+        at::Tensor q, at::Tensor k, at::Tensor v,
+        at::Tensor g_exp, at::Tensor beta, at::Tensor lengths
+    ) {
+        TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(),
+            "gated delta CUDA path requires CUDA tensors");
+        TORCH_CHECK(q.scalar_type() == at::kFloat && k.scalar_type() == at::kFloat &&
+                    v.scalar_type() == at::kFloat && g_exp.scalar_type() == at::kFloat &&
+                    beta.scalar_type() == at::kFloat,
+            "gated delta CUDA path expects FP32 working tensors");
+        const int64_t bh = q.size(0), seq = q.size(1);
+        const int64_t key_dim = q.size(2), val_dim = v.size(2);
+        int64_t heads_per_batch = 1;
+        const bool has_lengths = lengths.defined() && lengths.numel() > 0;
+        if (has_lengths) {
+            TORCH_CHECK(lengths.is_cuda() && lengths.device() == q.device() &&
+                        lengths.scalar_type() == at::kInt &&
+                        lengths.dim() == 1 && lengths.is_contiguous() &&
+                        lengths.size(0) > 0 && bh % lengths.size(0) == 0,
+                "gated delta lengths must be contiguous CUDA int32 [batch]");
+            heads_per_batch = bh / lengths.size(0);
+        }
+        auto state = at::zeros({bh, key_dim, val_dim}, q.options());
+        auto out = at::empty({bh, seq, val_dim}, q.options());
+        auto delta_buf = at::empty({bh, seq, val_dim}, q.options());
+        const int64_t checkpoint_stride =
+            gdn_state_checkpoint_stride(seq);
+        TORCH_CHECK(!env_enabled("QWEN36_GDN_CHUNKWISE_BWD") ||
+                checkpoint_stride > 0,
+            "QWEN36_GDN_CHUNKWISE_BWD requires a nonzero "
+            "QWEN36_GDN_STATE_CHECKPOINT_STRIDE no greater than the sequence length");
+        const int64_t checkpoint_count = checkpoint_stride > 0
+            ? (seq + checkpoint_stride - 1) / checkpoint_stride + 1
+            : 0;
+        auto state_checkpoints = checkpoint_stride > 0
+            ? at::empty(
+                {bh, checkpoint_count, key_dim, val_dim}, q.options())
+            : at::empty({0}, q.options());
+        auto stream = c10::cuda::getCurrentCUDAStream(q.device().index()).stream();
+        int status = cuda_gated_delta_rule(
+            q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(),
+            g_exp.data_ptr<float>(), beta.data_ptr<float>(),
+            state.data_ptr<float>(), out.data_ptr<float>(), delta_buf.data_ptr<float>(),
+            checkpoint_stride > 0
+                ? state_checkpoints.data_ptr<float>() : nullptr,
+            static_cast<int>(checkpoint_stride),
+            has_lengths ? lengths.data_ptr<int32_t>() : nullptr,
+            (int)heads_per_batch,
+            (int)bh, (int)seq, (int)key_dim, (int)val_dim,
+            reinterpret_cast<void*>(stream));
+        TORCH_CHECK(status == 0, "gated delta CUDA launch failed: ", status);
+        ctx->save_for_backward({
+            q, k, v, g_exp, beta, state, delta_buf, lengths,
+            state_checkpoints});
+        ctx->saved_data["checkpoint_stride"] = checkpoint_stride;
+        return out;
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output
+    ) {
+        auto saved = ctx->get_saved_variables();
+        auto q = saved[0];
+        auto k = saved[1];
+        auto v = saved[2];
+        auto g = saved[3];
+        auto beta = saved[4];
+        auto lengths = saved[7];
+        auto state_checkpoints = saved[8];
+        const int64_t checkpoint_stride =
+            ctx->saved_data["checkpoint_stride"].toInt();
+        const bool has_lengths = lengths.defined() && lengths.numel() > 0;
+        const int64_t heads_per_batch = has_lengths
+            ? q.size(0) / lengths.size(0)
+            : 1;
+        if (env_enabled("QWEN36_DELTA_REFERENCE_BWD")) {
+            auto q_ref = q.detach().set_requires_grad(true);
+            auto k_ref = k.detach().set_requires_grad(true);
+            auto v_ref = v.detach().set_requires_grad(true);
+            auto g_ref = g.detach().set_requires_grad(true);
+            auto beta_ref = beta.detach().set_requires_grad(true);
+            at::AutoGradMode guard(true);
+            auto reference = gated_delta_rule_reference(
+                q_ref, k_ref, v_ref, g_ref, beta_ref,
+                has_lengths ? lengths : at::Tensor(), heads_per_batch);
+            auto grads = torch::autograd::grad(
+                {reference}, {q_ref, k_ref, v_ref, g_ref, beta_ref},
+                {grad_output[0]}, /*retain_graph=*/false,
+                /*create_graph=*/false, /*allow_unused=*/false);
+            return {grads[0], grads[1], grads[2], grads[3], grads[4], at::Tensor()};
+        }
+
+        const int64_t bh = q.size(0), seq = q.size(1);
+        const int64_t key_dim = q.size(2), val_dim = v.size(2);
+        auto grad_out = grad_output[0].contiguous();
+        auto grad_q = at::empty_like(q);
+        auto grad_k = at::empty_like(k);
+        auto grad_v = at::empty_like(v);
+        auto grad_g = at::empty_like(g);
+        auto grad_beta = at::empty_like(beta);
+        auto stream = c10::cuda::getCurrentCUDAStream(q.device().index()).stream();
+        int status = cuda_gated_delta_rule_backward(
+            q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(),
+            g.data_ptr<float>(), beta.data_ptr<float>(),
+            saved[5].data_ptr<float>(), saved[6].data_ptr<float>(),
+            checkpoint_stride > 0
+                ? state_checkpoints.data_ptr<float>() : nullptr,
+            static_cast<int>(checkpoint_stride),
+            grad_out.data_ptr<float>(), grad_q.data_ptr<float>(),
+            grad_k.data_ptr<float>(), grad_v.data_ptr<float>(),
+            grad_g.data_ptr<float>(), grad_beta.data_ptr<float>(),
+            has_lengths ? lengths.data_ptr<int32_t>() : nullptr,
+            (int)heads_per_batch,
+            (int)bh, (int)seq, (int)key_dim, (int)val_dim,
+            reinterpret_cast<void*>(stream));
+        TORCH_CHECK(status == 0, "gated delta backward launch failed: ", status);
+        return {grad_q, grad_k, grad_v, grad_g, grad_beta, at::Tensor()};
+    }
+};
+
 static at::Tensor linear_attention(
     const at::Tensor& hidden,
     const at::Tensor& in_proj_qkv, const at::Tensor& in_proj_z,
@@ -378,7 +1495,8 @@ static at::Tensor linear_attention(
     int64_t num_k_heads, int64_t key_dim,
     int64_t num_v_heads, int64_t val_dim,
     int64_t conv_kernel, double rms_eps,
-    at::ScalarType compute_type
+    at::ScalarType compute_type,
+    const at::Tensor& attention_lengths
 ) {
     auto device = hidden.device();
     int64_t batch = hidden.size(0), seq = hidden.size(1);
@@ -391,7 +1509,18 @@ static at::Tensor linear_attention(
     int64_t seq_chunk = 0;
     if (chunk_env) seq_chunk = atoll(chunk_env);
 
-    if (seq_chunk > 0 && seq > seq_chunk) {
+    if (seq_chunk > 0 && seq > seq_chunk && conv_kernel > 1 &&
+        seq_chunk < conv_kernel - 1) {
+        TORCH_CHECK(false,
+            "QWEN36_SEQ_CHUNK must be at least conv_kernel - 1 for causal "
+            "overlap: seq_chunk=", seq_chunk,
+            " conv_kernel=", conv_kernel);
+    }
+
+    // Stateful chunking currently has no autograd state input/output contract.
+    // Keep it for inference/eval only; training uses the autograd-wrapped full
+    // sequence path below until chunk state gradients are implemented.
+    if (seq_chunk > 0 && seq > seq_chunk && !at::GradMode::is_enabled()) {
         // Chunked linear attention — mathematically equivalent to full sequence
         // Process in chunks, passing the delta rule state between chunks.
         // This avoids creating [batch, seq, qkv_dim] intermediate tensors.
@@ -450,8 +1579,10 @@ static at::Tensor linear_attention(
             q = q.repeat_interleave(n_rep, 2);
             k = k.repeat_interleave(n_rep, 2);
 
-            q = (q.to(at::kFloat) / q.to(at::kFloat).norm(2, -1, true).clamp_min(1e-6));
-            k = (k.to(at::kFloat) / k.to(at::kFloat).norm(2, -1, true).clamp_min(1e-6));
+            auto q_f = q.to(at::kFloat);
+            auto k_f = k.to(at::kFloat);
+            q = q_f * (q_f.pow(2).sum(-1, true) + 1e-6).rsqrt();
+            k = k_f * (k_f.pow(2).sum(-1, true) + 1e-6).rsqrt();
             q = q * (1.0 / std::sqrt((double)key_dim));
 
             auto q_t = q.transpose(1, 2).contiguous();
@@ -469,9 +1600,13 @@ static at::Tensor linear_attention(
             auto state_contig = state.contiguous();
             auto outs = at::empty({BH, chunk_len, val_dim}, q_t.options());
             auto delta_buf = at::empty({BH, chunk_len, val_dim}, q_t.options());
+            auto chunk_lengths = attention_lengths.defined()
+                ? (attention_lengths - offset).clamp(0, chunk_len).to(at::kInt).contiguous()
+                : at::Tensor();
 
             // CUDA kernel — state is passed in and updated in-place
-            cuda_gated_delta_rule(
+            auto stream = c10::cuda::getCurrentCUDAStream(device.index()).stream();
+            int status = cuda_gated_delta_rule(
                 q_contig.data_ptr<float>(),
                 k_contig.data_ptr<float>(),
                 v_contig.data_ptr<float>(),
@@ -480,8 +1615,13 @@ static at::Tensor linear_attention(
                 state_contig.data_ptr<float>(),
                 outs.data_ptr<float>(),
                 delta_buf.data_ptr<float>(),
-                (int)BH, (int)chunk_len, (int)key_dim, (int)val_dim
+                nullptr, 0,
+                chunk_lengths.defined() ? chunk_lengths.data_ptr<int32_t>() : nullptr,
+                (int)num_v_heads,
+                (int)BH, (int)chunk_len, (int)key_dim, (int)val_dim,
+                reinterpret_cast<void*>(stream)
             );
+            TORCH_CHECK(status == 0, "gated delta CUDA launch failed: ", status);
             state = state_contig;  // updated state for next chunk
 
             auto core_out = outs.reshape({batch, num_v_heads, chunk_len, val_dim})
@@ -545,8 +1685,10 @@ static at::Tensor linear_attention(
     k = k.repeat_interleave(n_rep, 2);
 
     // L2 normalize Q, K (HF: use_qk_l2norm_in_kernel=True, eps=1e-6)
-    q = (q.to(at::kFloat) / q.to(at::kFloat).norm(2, -1, true).clamp_min(1e-6));
-    k = (k.to(at::kFloat) / k.to(at::kFloat).norm(2, -1, true).clamp_min(1e-6));
+    auto q_f = q.to(at::kFloat);
+    auto k_f = k.to(at::kFloat);
+    q = q_f * (q_f.pow(2).sum(-1, true) + 1e-6).rsqrt();
+    k = k_f * (k_f.pow(2).sum(-1, true) + 1e-6).rsqrt();
 
     // Scale Q by 1/sqrt(key_dim) — matching HF: scale = 1 / (key_dim ** 0.5)
     double scale = 1.0 / std::sqrt((double)key_dim);
@@ -563,30 +1705,15 @@ static at::Tensor linear_attention(
     // All computation runs in a single kernel launch using shared memory for state
     auto g_exp = g_t.exp();  // [B, H, S]
     int64_t BH = batch * num_v_heads;
-    auto state = at::zeros({BH, key_dim, val_dim}, q_t.options());  // [B*H, D_k, D_v]
-
-    // Prepare contiguous FP32 tensors for CUDA kernel
+    // Prepare contiguous FP32 tensors for the CUDA forward/autograd wrapper.
     auto q_contig = q_t.reshape({BH, seq, key_dim}).contiguous().to(at::kFloat);
     auto k_contig = k_t.reshape({BH, seq, key_dim}).contiguous().to(at::kFloat);
     auto v_contig = v_t.reshape({BH, seq, val_dim}).contiguous().to(at::kFloat);
     auto g_contig = g_exp.reshape({BH, seq}).contiguous().to(at::kFloat);
     auto beta_contig = beta_t.reshape({BH, seq}).contiguous().to(at::kFloat);
-    auto state_contig = state.contiguous();
-    auto outs = at::empty({BH, seq, val_dim}, q_t.options());
-    auto delta_buf = at::empty({BH, seq, val_dim}, q_t.options());
-
-    // Launch CUDA kernel — single launch replaces seq×3 bmm calls
-    cuda_gated_delta_rule(
-        q_contig.data_ptr<float>(),
-        k_contig.data_ptr<float>(),
-        v_contig.data_ptr<float>(),
-        g_contig.data_ptr<float>(),
-        beta_contig.data_ptr<float>(),
-        state_contig.data_ptr<float>(),
-        outs.data_ptr<float>(),
-        delta_buf.data_ptr<float>(),
-        (int)BH, (int)seq, (int)key_dim, (int)val_dim
-    );
+    auto outs = GatedDeltaRuleFunction::apply(
+        q_contig, k_contig, v_contig, g_contig, beta_contig,
+        gdn_lengths_arg(attention_lengths, q_contig));
 
     // Reshape: [B*H, S, D_v] → [B, H, S, D_v] → [B, S, H, D_v]
     auto core_out = outs.reshape({batch, num_v_heads, seq, val_dim})
@@ -628,64 +1755,1471 @@ static at::Tensor full_attention_batched(TrainingContext* ctx, const at::Tensor&
     const at::Tensor& v_proj, const at::Tensor& o_proj,
     int64_t num_heads, int64_t num_kv_heads, int64_t head_dim,
     double partial_rotary_factor, double rope_theta,
-    double rms_eps, at::ScalarType kind, const at::Tensor& attention_mask);
+    double rms_eps, at::ScalarType kind, const at::Tensor& attention_mask,
+    const at::Tensor& fused_qkv);
 static at::Tensor linear_attention_batched(TrainingContext* ctx, const at::Tensor& hidden,
     int64_t layer_idx, const at::Tensor& in_proj_qkv, const at::Tensor& in_proj_z,
     const at::Tensor& in_proj_a, const at::Tensor& in_proj_b,
     const at::Tensor& a_log, const at::Tensor& dt_bias,
     const at::Tensor& conv1d_w, const at::Tensor& norm_w, const at::Tensor& out_proj,
     int64_t num_k_heads, int64_t key_dim, int64_t num_v_heads, int64_t val_dim,
-    int64_t conv_kernel, double rms_eps, at::ScalarType compute_type);
+    int64_t conv_kernel, double rms_eps, at::ScalarType compute_type,
+    const at::Tensor& attention_lengths);
 
 // ──────────────────────────────────────────────────────────────────────
 // MoE
 // ──────────────────────────────────────────────────────────────────────
 
-static at::Tensor moe_forward(
-    void* nccl_comm_v, void* nccl_stream_v,
-    const at::Tensor& hidden,
-    const at::Tensor& gate_w, const at::Tensor& shared_expert_gate_w,
-    const at::Tensor& shared_gate_proj, const at::Tensor& shared_up_proj, const at::Tensor& shared_down_proj,
-    const at::Tensor& experts_gate_up, const at::Tensor& experts_down,
-    int64_t num_experts, int64_t top_k, int64_t intermediate,
-    bool norm_topk_prob, int64_t expert_start, int64_t expert_count,
-    at::ScalarType compute_type
-) {
-    int64_t batch = hidden.size(0), seq = hidden.size(1), hidden_dim = hidden.size(2);
-    auto device = hidden.device();
-    auto flat = hidden.reshape({batch * seq, hidden_dim});
-    int64_t N = flat.size(0);
+static ncclDataType_t qwen36_nccl_dtype(at::ScalarType type) {
+    switch (type) {
+        case at::kFloat: return ncclFloat32;
+        case at::kHalf: return ncclFloat16;
+        case at::kBFloat16: return ncclBfloat16;
+        case at::kInt: return ncclInt;
+        case at::kLong: return ncclInt64;
+        default:
+            TORCH_CHECK(false, "unsupported NCCL dtype: ", type);
+    }
+}
 
-    auto router_logits = at::matmul(flat, gate_w.t());
-    auto routing_weights = router_logits.softmax(-1, at::kFloat);
-    auto [topk_weights, topk_indices] = routing_weights.topk(top_k, -1, true, true);
-    if (norm_topk_prob) {
-        auto denom = topk_weights.sum(-1, true).clamp_min(1e-9);
-        topk_weights = topk_weights / denom;
+struct Qwen36A2ACountPlan {
+    at::Tensor send_offsets;
+    at::Tensor receive_offsets;
+};
+
+struct Qwen36A2ACountExchange {
+    at::Tensor all_counts;
+    at::Tensor staged_counts;
+    at::Tensor host_counts;
+    cudaStream_t metadata_stream = nullptr;
+    cudaEvent_t ready = nullptr;
+    cudaEvent_t done = nullptr;
+    int device = -1;
+    int world = 0;
+    int rank = 0;
+    bool compact_metadata = false;
+    bool completion_recorded = false;
+
+    Qwen36A2ACountExchange() = default;
+    Qwen36A2ACountExchange(const Qwen36A2ACountExchange&) = delete;
+    Qwen36A2ACountExchange& operator=(const Qwen36A2ACountExchange&) = delete;
+
+    Qwen36A2ACountExchange(Qwen36A2ACountExchange&& other) noexcept
+        : all_counts(std::move(other.all_counts)),
+          staged_counts(std::move(other.staged_counts)),
+          host_counts(std::move(other.host_counts)),
+          metadata_stream(other.metadata_stream),
+          ready(other.ready),
+          done(other.done),
+          device(other.device),
+          world(other.world),
+          rank(other.rank),
+          compact_metadata(other.compact_metadata),
+          completion_recorded(other.completion_recorded) {
+        other.metadata_stream = nullptr;
+        other.ready = nullptr;
+        other.done = nullptr;
+        other.device = -1;
+        other.completion_recorded = false;
     }
-    topk_weights = topk_weights.to(compute_type);
 
-    auto routed_output = at::zeros(flat.sizes(), flat.options());
+    ~Qwen36A2ACountExchange() {
+        // A packing/allocation exception after launch must not let the caching
+        // allocator recycle tensors while NCCL or the D2H copy still uses them.
+        int previous_device = -1;
+        const bool restore_device = device >= 0 &&
+            cudaGetDevice(&previous_device) == cudaSuccess &&
+            previous_device != device &&
+            cudaSetDevice(device) == cudaSuccess;
+        if (completion_recorded && done) {
+            cudaEventSynchronize(done);
+        } else if (metadata_stream) {
+            cudaStreamSynchronize(metadata_stream);
+        }
+        if (ready) cudaEventDestroy(ready);
+        if (done) cudaEventDestroy(done);
+        if (restore_device) cudaSetDevice(previous_device);
+    }
 
-    // Debug: dump MoE routing and weight stats
-    if (getenv("QWEN36_DUMP_MOE")) {
-        auto rl_f = router_logits.to(at::kFloat);
-        auto rw_f = topk_weights.to(at::kFloat);
-        auto egu_f = experts_gate_up.select(0, 0).to(at::kFloat);
-        auto ed_f = experts_down.select(0, 0).to(at::kFloat);
-        auto sg_f = shared_gate_proj.to(at::kFloat);
-        auto sd_f = shared_down_proj.to(at::kFloat);
-        auto seg_f = at::sigmoid(at::matmul(flat, shared_expert_gate_w.t())).to(at::kFloat);
+    void wait() {
+        if (!completion_recorded) return;
+        int previous_device = -1;
+        const bool restore_device = device >= 0 &&
+            cudaGetDevice(&previous_device) == cudaSuccess &&
+            previous_device != device &&
+            cudaSetDevice(device) == cudaSuccess;
+        const auto status = cudaEventSynchronize(done);
+        if (restore_device) cudaSetDevice(previous_device);
+        TORCH_CHECK(status == cudaSuccess,
+            "failed to synchronize EP A2A count metadata");
+        completion_recorded = false;
+        metadata_stream = nullptr;
     }
+};
 
-    // Sort-based expert dispatch — eliminates eq/nonzero/index_select per expert.
-    // At large N (1000+), the old for-loop with eq+nonzero per expert was O(experts × tokens)
-    // with high kernel launch overhead. This approach pre-sorts tokens by expert assignment.
-    //
-    // Autograd note: routing indices/weights are computed in no-grad forward (detached).
-    // Only matmul inputs/outputs participate in autograd. sort/index_select/index_add
+static Qwen36A2ACountExchange qwen36_a2a_counts_launch(
+    const at::Tensor& local_metadata, ncclComm_t comm,
+    cudaStream_t current_stream, cudaStream_t metadata_stream,
+    bool overlap
+) {
+    Qwen36A2ACountExchange exchange;
+    int world = 0;
+    auto err = ncclCommCount(comm, &world);
+    TORCH_CHECK(err == ncclSuccess, "ncclCommCount failed: ",
+        ncclGetErrorString(err));
+    TORCH_CHECK(local_metadata.numel() == world + 1,
+        "EP A2A metadata must contain one count per peer and a validity flag");
+    exchange.device = local_metadata.device().index();
+    exchange.world = world;
+    exchange.all_counts = at::empty(
+        {world, world + 1}, local_metadata.options());
+    const bool use_overlap = overlap && metadata_stream &&
+        metadata_stream != current_stream;
+    auto count_stream = current_stream;
+    if (use_overlap) {
+        exchange.metadata_stream = metadata_stream;
+        count_stream = metadata_stream;
+        TORCH_CHECK(cudaEventCreateWithFlags(
+                &exchange.ready, cudaEventDisableTiming) == cudaSuccess,
+            "failed to create EP A2A count producer event");
+        TORCH_CHECK(cudaEventCreateWithFlags(
+                &exchange.done, cudaEventDisableTiming) == cudaSuccess,
+            "failed to create EP A2A count completion event");
+        TORCH_CHECK(cudaEventRecord(exchange.ready, current_stream) == cudaSuccess,
+            "failed to record EP A2A count producer event");
+        TORCH_CHECK(cudaStreamWaitEvent(
+                metadata_stream, exchange.ready, 0) == cudaSuccess,
+            "failed to fence EP A2A metadata stream");
+    }
+    err = ncclAllGather(
+        local_metadata.data_ptr(), exchange.all_counts.data_ptr(),
+        world + 1, ncclInt, comm, count_stream);
+    TORCH_CHECK(err == ncclSuccess, "EP A2A count all-gather failed: ",
+        ncclGetErrorString(err));
+    int rank = 0;
+    err = ncclCommUserRank(comm, &rank);
+    TORCH_CHECK(err == ncclSuccess, "ncclCommUserRank failed: ",
+        ncclGetErrorString(err));
+    exchange.rank = rank;
+
+    // The legacy path copies the complete world-by-world metadata matrix to
+    // CPU. For larger EP groups, only this rank's send row, receive column,
+    // and one validity flag per peer are needed by the host-side NCCL
+    // variable-count enqueue. Keep the compact path opt-in for small worlds
+    // because its extra device gather is not worth it at EP2/EP4.
+    exchange.compact_metadata = env_enabled(
+        "QWEN36_EP_A2A_COMPACT_COUNTS", world >= 8);
+    if (use_overlap) {
+
+        const auto side_stream = c10::cuda::getStreamFromExternal(
+            metadata_stream, local_metadata.device().index());
+        {
+            c10::cuda::CUDAStreamGuard guard(side_stream);
+            if (exchange.compact_metadata) {
+                auto send_row = exchange.all_counts.select(0, rank)
+                    .narrow(0, 0, world);
+                auto receive_column = exchange.all_counts.narrow(1, rank, 1)
+                    .select(1, 0).narrow(0, 0, world);
+                auto invalid_flags = exchange.all_counts.select(1, world);
+                exchange.staged_counts = at::cat(
+                    {send_row, receive_column, invalid_flags}, 0).contiguous();
+            } else {
+                exchange.staged_counts = exchange.all_counts;
+            }
+        }
+        exchange.host_counts = at::empty(
+            exchange.staged_counts.sizes(),
+            at::TensorOptions().device(at::kCPU).dtype(at::kInt)
+                .pinned_memory(true));
+        TORCH_CHECK(cudaMemcpyAsync(
+                exchange.host_counts.data_ptr<int32_t>(),
+                exchange.staged_counts.data_ptr<int32_t>(),
+                exchange.staged_counts.numel() * sizeof(int32_t),
+                cudaMemcpyDeviceToHost, metadata_stream) == cudaSuccess,
+            "failed to copy EP A2A count metadata to pinned host memory");
+        TORCH_CHECK(cudaEventRecord(exchange.done, metadata_stream) == cudaSuccess,
+            "failed to record EP A2A count completion event");
+        exchange.completion_recorded = true;
+    } else {
+        if (exchange.compact_metadata) {
+            auto send_row = exchange.all_counts.select(0, rank)
+                .narrow(0, 0, world);
+            auto receive_column = exchange.all_counts.narrow(1, rank, 1)
+                .select(1, 0).narrow(0, 0, world);
+            auto invalid_flags = exchange.all_counts.select(1, world);
+            exchange.host_counts = at::cat(
+                {send_row, receive_column, invalid_flags}, 0)
+                .contiguous()
+                .to(at::TensorOptions().device(at::kCPU));
+        } else {
+            exchange.host_counts = exchange.all_counts.to(
+                at::TensorOptions().device(at::kCPU));
+        }
+    }
+    return exchange;
+}
+
+static Qwen36A2ACountPlan qwen36_a2a_counts_finalize(
+    Qwen36A2ACountExchange& exchange
+) {
+    exchange.wait();
+    const int world = exchange.world;
+    const int rank = exchange.rank;
+    const bool compact_metadata = exchange.compact_metadata;
+    const auto& host = exchange.host_counts;
+    Qwen36A2ACountPlan plan;
+    // Dispatch forward, dispatch backward, combine forward, and combine
+    // backward all use the same variable-count layout. Materialize the two
+    // prefix plans once and keep them in the autograd graph instead of
+    // repeatedly copying count tensors into vectors and rescanning them.
+    auto offset_storage = at::empty(
+        {2, world + 1},
+        at::TensorOptions().device(at::kCPU).dtype(at::kLong));
+    plan.send_offsets = offset_storage.select(0, 0);
+    plan.receive_offsets = offset_storage.select(0, 1);
+    auto* send_offsets = plan.send_offsets.data_ptr<int64_t>();
+    auto* receive_offsets = plan.receive_offsets.data_ptr<int64_t>();
+    send_offsets[0] = 0;
+    receive_offsets[0] = 0;
+    for (int peer = 0; peer < world; ++peer) {
+        const int32_t send_count = compact_metadata
+            ? host.data_ptr<int32_t>()[peer]
+            : host.data_ptr<int32_t>()[rank * (world + 1) + peer];
+        const int32_t receive_count = compact_metadata
+            ? host.data_ptr<int32_t>()[world + peer]
+            : host.data_ptr<int32_t>()[peer * (world + 1) + rank];
+        TORCH_CHECK(send_count >= 0 && receive_count >= 0,
+            "negative EP A2A token count");
+        const int32_t invalid = compact_metadata
+            ? host.data_ptr<int32_t>()[2 * world + peer]
+            : host.data_ptr<int32_t>()[peer * (world + 1) + world];
+        TORCH_CHECK(invalid == 0,
+            "EP A2A expert index is outside the communicator range");
+        send_offsets[peer + 1] = send_offsets[peer] + send_count;
+        receive_offsets[peer + 1] = receive_offsets[peer] + receive_count;
+    }
+    return plan;
+}
+
+// Variable-split token dispatch. Metadata is deliberately non-differentiable;
+// only the hidden activation participates in the custom backward exchange.
+struct Qwen36A2ADispatchFunction : public torch::autograd::Function<Qwen36A2ADispatchFunction> {
+    static std::vector<at::Tensor> forward(
+        torch::autograd::AutogradContext* ctx,
+        at::Tensor input, at::Tensor expert_indices, at::Tensor token_indices,
+        int64_t expert_count, int64_t comm_ptr, int64_t stream_ptr,
+        int64_t metadata_stream_ptr
+    ) {
+        auto comm = reinterpret_cast<ncclComm_t>(comm_ptr);
+        TORCH_CHECK(comm, "EP A2A requires an NCCL communicator");
+        int world = 1, rank = 0;
+        TORCH_CHECK(ncclCommCount(comm, &world) == ncclSuccess,
+            "ncclCommCount failed");
+        TORCH_CHECK(ncclCommUserRank(comm, &rank) == ncclSuccess,
+            "ncclCommUserRank failed");
+        TORCH_CHECK(expert_count > 0 && expert_indices.scalar_type() == at::kLong,
+            "invalid EP A2A expert metadata");
+        TORCH_CHECK(expert_indices.numel() == token_indices.numel(),
+            "EP A2A expert and source-token metadata must have equal length");
+        auto stream = c10::cuda::getCurrentCUDAStream(input.device().index()).stream();
+        const int64_t hidden = input.size(1);
+        auto gpu_count_opts = at::TensorOptions()
+            .device(input.device()).dtype(at::kInt);
+        at::Tensor send_index;
+        at::Tensor local_metadata;
+        // Peer-wise nonzero is cheaper at EP2; one GPU sort wins as the EP
+        // fan-out grows for sufficiently large token batches. The override
+        // keeps both paths available for A/B and production tuning.
+        const bool use_gpu_metadata = env_enabled(
+            "QWEN36_EP_A2A_GPU_METADATA",
+            world >= 4 && expert_indices.numel() >= 512);
+        if (use_gpu_metadata) {
+            auto destinations = at::floor_divide(
+                expert_indices, expert_count);
+            auto valid_destinations = destinations.clamp(0, world - 1);
+            auto invalid = destinations.ne(valid_destinations)
+                .any().to(at::kInt).reshape({1});
+            auto [sorted_destinations, order] = valid_destinations.sort(0);
+            send_index = std::move(order);
+            auto local_counts = at::bincount(
+                sorted_destinations, c10::nullopt, world).to(at::kInt);
+            local_metadata = at::cat({local_counts, invalid}, 0);
+        } else {
+            std::vector<at::Tensor> indices(world);
+            std::vector<int32_t> host_metadata(world + 1, 0);
+            for (int dst = 0; dst < world; ++dst) {
+                auto mask = (expert_indices >= dst * expert_count) &
+                    (expert_indices < (dst + 1) * expert_count);
+                indices[dst] = at::nonzero(mask).reshape({-1});
+                host_metadata[dst] = static_cast<int32_t>(
+                    indices[dst].numel());
+            }
+            send_index = at::cat(indices, 0);
+            host_metadata[world] = send_index.numel() != expert_indices.numel();
+            auto host_counts = at::from_blob(
+                host_metadata.data(), {world + 1},
+                at::TensorOptions().device(at::kCPU).dtype(at::kInt));
+            local_metadata = host_counts.to(gpu_count_opts);
+        }
+        const auto metadata_stream =
+            reinterpret_cast<cudaStream_t>(metadata_stream_ptr);
+        const bool count_overlap =
+            env_enabled("QWEN36_EP_A2A_COUNT_OVERLAP") &&
+            metadata_stream && metadata_stream != stream;
+        auto count_exchange = qwen36_a2a_counts_launch(
+            local_metadata, comm, stream,
+            metadata_stream, count_overlap);
+        Qwen36A2ACountPlan count_plan;
+        if (!count_overlap) {
+            // Preserve the legacy validation/error boundary before payload
+            // packing when overlap is disabled or unavailable.
+            count_plan = qwen36_a2a_counts_finalize(count_exchange);
+        }
+        auto send_token = token_indices.index_select(0, send_index).contiguous();
+        auto send_hidden = input.index_select(0, send_token).contiguous();
+        auto send_expert = expert_indices.index_select(0, send_index).contiguous();
+        // Token and expert IDs always travel together and are both int64.
+        // Packing them removes one NCCL enqueue per peer without changing the
+        // variable-count protocol or the autograd payload. Keep the legacy
+        // split form as a runtime fallback for protocol A/B and rollback.
+        const bool packed_metadata = env_enabled(
+            "QWEN36_EP_A2A_PACKED_METADATA", true);
+        at::Tensor send_metadata;
+        if (packed_metadata) {
+            send_metadata = at::stack(
+                {send_token, send_expert}, 1).contiguous();
+        }
+        if (count_overlap) {
+            count_plan = qwen36_a2a_counts_finalize(count_exchange);
+        }
+        auto send_offsets_tensor = std::move(count_plan.send_offsets);
+        auto recv_offsets_tensor = std::move(count_plan.receive_offsets);
+        const auto* send_offsets = send_offsets_tensor.data_ptr<int64_t>();
+        const auto* recv_offsets = recv_offsets_tensor.data_ptr<int64_t>();
+        auto recv_hidden = at::empty({recv_offsets[world], hidden}, input.options());
+        at::Tensor recv_metadata;
+        at::Tensor recv_token;
+        at::Tensor recv_expert;
+        if (packed_metadata) {
+            recv_metadata = at::empty(
+                {recv_offsets[world], 2}, token_indices.options());
+            recv_token = recv_metadata.select(1, 0);
+            recv_expert = recv_metadata.select(1, 1);
+        } else {
+            recv_token = at::empty({recv_offsets[world]}, token_indices.options());
+            recv_expert = at::empty({recv_offsets[world]}, expert_indices.options());
+        }
+        auto send_ptr = [&](at::Tensor& t, int64_t row, int64_t width) -> const void* {
+            return static_cast<const char*>(t.data_ptr()) +
+                row * width * t.element_size();
+        };
+        auto recv_ptr = [&](at::Tensor& t, int64_t row, int64_t width) -> void* {
+            return static_cast<char*>(t.data_ptr()) +
+                row * width * t.element_size();
+        };
+        Qwen36StreamFence stream_fence(
+            stream, reinterpret_cast<cudaStream_t>(stream_ptr));
+        auto operation_stream = stream_fence.operation;
+        TORCH_CHECK(ncclGroupStart() == ncclSuccess, "ncclGroupStart failed");
+        for (int peer = 0; peer < world; ++peer) {
+            const int64_t rows = send_offsets[peer + 1] - send_offsets[peer];
+            if (rows) {
+                auto err = ncclSend(send_ptr(send_hidden, send_offsets[peer], hidden),
+                    rows * hidden, qwen36_nccl_dtype(input.scalar_type()), peer, comm, operation_stream);
+                TORCH_CHECK(err == ncclSuccess, "A2A hidden send failed: ", ncclGetErrorString(err));
+                if (packed_metadata) {
+                    err = ncclSend(
+                        static_cast<const char*>(send_metadata.data_ptr()) +
+                            send_offsets[peer] * 2 * send_metadata.element_size(),
+                        rows * 2, ncclInt64, peer, comm, operation_stream);
+                    TORCH_CHECK(err == ncclSuccess,
+                        "A2A token/expert metadata send failed: ",
+                        ncclGetErrorString(err));
+                } else {
+                    err = ncclSend(
+                        send_ptr(send_token, send_offsets[peer], 1), rows,
+                        ncclInt64, peer, comm, operation_stream);
+                    TORCH_CHECK(err == ncclSuccess,
+                        "A2A token send failed: ", ncclGetErrorString(err));
+                    err = ncclSend(
+                        send_ptr(send_expert, send_offsets[peer], 1), rows,
+                        ncclInt64, peer, comm, operation_stream);
+                    TORCH_CHECK(err == ncclSuccess,
+                        "A2A expert send failed: ", ncclGetErrorString(err));
+                }
+            }
+        }
+        for (int peer = 0; peer < world; ++peer) {
+            const int64_t rows = recv_offsets[peer + 1] - recv_offsets[peer];
+            if (rows) {
+                auto err = ncclRecv(recv_ptr(recv_hidden, recv_offsets[peer], hidden),
+                    rows * hidden, qwen36_nccl_dtype(input.scalar_type()), peer, comm, operation_stream);
+                TORCH_CHECK(err == ncclSuccess, "A2A hidden recv failed: ", ncclGetErrorString(err));
+                if (packed_metadata) {
+                    err = ncclRecv(
+                        static_cast<char*>(recv_metadata.data_ptr()) +
+                            recv_offsets[peer] * 2 * recv_metadata.element_size(),
+                        rows * 2, ncclInt64, peer, comm, operation_stream);
+                    TORCH_CHECK(err == ncclSuccess,
+                        "A2A token/expert metadata recv failed: ",
+                        ncclGetErrorString(err));
+                } else {
+                    err = ncclRecv(
+                        recv_ptr(recv_token, recv_offsets[peer], 1), rows,
+                        ncclInt64, peer, comm, operation_stream);
+                    TORCH_CHECK(err == ncclSuccess,
+                        "A2A token recv failed: ", ncclGetErrorString(err));
+                    err = ncclRecv(
+                        recv_ptr(recv_expert, recv_offsets[peer], 1), rows,
+                        ncclInt64, peer, comm, operation_stream);
+                    TORCH_CHECK(err == ncclSuccess,
+                        "A2A expert recv failed: ", ncclGetErrorString(err));
+                }
+            }
+        }
+        TORCH_CHECK(ncclGroupEnd() == ncclSuccess, "A2A dispatch group failed");
+        stream_fence.complete();
+        auto recv_local = recv_expert - rank * expert_count;
+        ctx->save_for_backward({input, send_token, send_offsets_tensor, recv_offsets_tensor});
+        ctx->saved_data["comm"] = comm_ptr;
+        ctx->saved_data["stream"] = stream_ptr;
+        ctx->saved_data["expert_count"] = expert_count;
+        return {recv_hidden, recv_token, recv_local, send_index,
+            send_offsets_tensor, recv_offsets_tensor};
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx, std::vector<at::Tensor> grad_output
+    ) {
+        auto saved = ctx->get_saved_variables();
+        auto input = saved[0];
+        auto send_token = saved[1];
+        const auto& send_offsets = saved[2];
+        const auto& recv_offsets = saved[3];
+        const int world = send_offsets.numel() - 1;
+        TORCH_CHECK(world >= 1 && recv_offsets.numel() == world + 1,
+            "invalid saved EP A2A prefix plan");
+        const auto* so = send_offsets.data_ptr<int64_t>();
+        const auto* ro = recv_offsets.data_ptr<int64_t>();
+        auto comm = reinterpret_cast<ncclComm_t>(ctx->saved_data["comm"].toInt());
+        auto stream = c10::cuda::getCurrentCUDAStream(input.device().index()).stream();
+        auto requested_stream = reinterpret_cast<cudaStream_t>(
+            ctx->saved_data["stream"].toInt());
+        auto grad_input = at::zeros_like(input);
+        auto returned = at::empty({so[world], input.size(1)}, input.options());
+        const size_t elem_bytes = input.element_size();
+        Qwen36StreamFence stream_fence(stream, requested_stream);
+        auto operation_stream = stream_fence.operation;
+        TORCH_CHECK(ncclGroupStart() == ncclSuccess, "ncclGroupStart failed");
+        for (int peer = 0; peer < world; ++peer) if (so[peer + 1] != so[peer]) {
+            const int64_t rows = so[peer + 1] - so[peer];
+            auto err = ncclRecv(static_cast<char*>(returned.data_ptr()) + so[peer] * input.size(1) * elem_bytes,
+                rows * input.size(1), qwen36_nccl_dtype(input.scalar_type()), peer, comm, operation_stream);
+            TORCH_CHECK(err == ncclSuccess, "A2A backward recv failed: ", ncclGetErrorString(err));
+        }
+        for (int peer = 0; peer < world; ++peer) if (ro[peer + 1] != ro[peer]) {
+            const int64_t rows = ro[peer + 1] - ro[peer];
+            auto err = ncclSend(static_cast<const char*>(grad_output[0].data_ptr()) + ro[peer] * input.size(1) * elem_bytes,
+                rows * input.size(1), qwen36_nccl_dtype(input.scalar_type()), peer, comm, operation_stream);
+            TORCH_CHECK(err == ncclSuccess, "A2A backward send failed: ", ncclGetErrorString(err));
+        }
+        TORCH_CHECK(ncclGroupEnd() == ncclSuccess, "A2A dispatch backward group failed");
+        stream_fence.complete();
+        grad_input.index_add_(0, send_token, returned);
+        return {grad_input, at::Tensor(), at::Tensor(), at::Tensor(),
+            at::Tensor(), at::Tensor(), at::Tensor()};
+    }
+};
+
+struct Qwen36A2ACombineFunction : public torch::autograd::Function<Qwen36A2ACombineFunction> {
+    static at::Tensor forward(torch::autograd::AutogradContext* ctx,
+        at::Tensor local_output, at::Tensor send_offsets,
+        at::Tensor recv_offsets, int64_t comm_ptr, int64_t stream_ptr) {
+        auto comm = reinterpret_cast<ncclComm_t>(comm_ptr);
+        const int world = send_offsets.numel() - 1;
+        TORCH_CHECK(world >= 1 && recv_offsets.numel() == world + 1,
+            "invalid EP A2A prefix plan");
+        const auto* so = send_offsets.data_ptr<int64_t>();
+        const auto* ro = recv_offsets.data_ptr<int64_t>();
+        auto returned = at::empty({so[world], local_output.size(1)}, local_output.options());
+        const size_t elem_bytes = local_output.element_size();
+        auto current_stream = c10::cuda::getCurrentCUDAStream(
+            local_output.device().index()).stream();
+        Qwen36StreamFence stream_fence(
+            current_stream, reinterpret_cast<cudaStream_t>(stream_ptr));
+        auto operation_stream = stream_fence.operation;
+        TORCH_CHECK(ncclGroupStart() == ncclSuccess, "ncclGroupStart failed");
+        for (int peer = 0; peer < world; ++peer) if (ro[peer + 1] != ro[peer]) {
+            const int64_t rows = ro[peer + 1] - ro[peer];
+            auto err = ncclSend(static_cast<const char*>(local_output.data_ptr()) + ro[peer] * local_output.size(1) * elem_bytes,
+                rows * local_output.size(1), qwen36_nccl_dtype(local_output.scalar_type()), peer, comm, operation_stream);
+            TORCH_CHECK(err == ncclSuccess, "A2A combine send failed: ", ncclGetErrorString(err));
+        }
+        for (int peer = 0; peer < world; ++peer) if (so[peer + 1] != so[peer]) {
+            const int64_t rows = so[peer + 1] - so[peer];
+            auto err = ncclRecv(static_cast<char*>(returned.data_ptr()) + so[peer] * local_output.size(1) * elem_bytes,
+                rows * local_output.size(1), qwen36_nccl_dtype(local_output.scalar_type()), peer, comm, operation_stream);
+            TORCH_CHECK(err == ncclSuccess, "A2A combine recv failed: ", ncclGetErrorString(err));
+        }
+        TORCH_CHECK(ncclGroupEnd() == ncclSuccess, "A2A combine group failed");
+        stream_fence.complete();
+        ctx->save_for_backward({send_offsets, recv_offsets});
+        ctx->saved_data["comm"] = comm_ptr;
+        ctx->saved_data["stream"] = stream_ptr;
+        return returned;
+    }
+
+    static std::vector<at::Tensor> backward(torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output) {
+        auto saved = ctx->get_saved_variables();
+        const auto& send_offsets = saved[0];
+        const auto& recv_offsets = saved[1];
+        const int world = send_offsets.numel() - 1;
+        TORCH_CHECK(world >= 1 && recv_offsets.numel() == world + 1,
+            "invalid saved EP A2A prefix plan");
+        const auto* so = send_offsets.data_ptr<int64_t>();
+        const auto* ro = recv_offsets.data_ptr<int64_t>();
+        auto comm = reinterpret_cast<ncclComm_t>(ctx->saved_data["comm"].toInt());
+        auto stream = c10::cuda::getCurrentCUDAStream(grad_output[0].device().index()).stream();
+        auto requested_stream = reinterpret_cast<cudaStream_t>(
+            ctx->saved_data["stream"].toInt());
+        auto packed = grad_output[0].contiguous();
+        auto grad_local = at::empty({ro[world], grad_output[0].size(1)}, grad_output[0].options());
+        const size_t elem_bytes = grad_output[0].element_size();
+        Qwen36StreamFence stream_fence(stream, requested_stream);
+        auto operation_stream = stream_fence.operation;
+        TORCH_CHECK(ncclGroupStart() == ncclSuccess, "ncclGroupStart failed");
+        for (int peer = 0; peer < world; ++peer) if (so[peer + 1] != so[peer]) {
+            const int64_t rows = so[peer + 1] - so[peer];
+            auto err = ncclSend(static_cast<const char*>(packed.data_ptr()) + so[peer] * grad_output[0].size(1) * elem_bytes,
+                rows * grad_output[0].size(1), qwen36_nccl_dtype(grad_output[0].scalar_type()), peer, comm, operation_stream);
+            TORCH_CHECK(err == ncclSuccess, "A2A combine backward send failed: ", ncclGetErrorString(err));
+        }
+        for (int peer = 0; peer < world; ++peer) if (ro[peer + 1] != ro[peer]) {
+            const int64_t rows = ro[peer + 1] - ro[peer];
+            auto err = ncclRecv(static_cast<char*>(grad_local.data_ptr()) + ro[peer] * grad_output[0].size(1) * elem_bytes,
+                rows * grad_output[0].size(1), qwen36_nccl_dtype(grad_output[0].scalar_type()), peer, comm, operation_stream);
+            TORCH_CHECK(err == ncclSuccess, "A2A combine backward recv failed: ", ncclGetErrorString(err));
+        }
+        TORCH_CHECK(ncclGroupEnd() == ncclSuccess, "A2A combine backward group failed");
+        stream_fence.complete();
+        return {grad_local, at::Tensor(), at::Tensor(), at::Tensor(),
+            at::Tensor()};
+    }
+};
+
+struct Qwen36FusedWeightedUnpermuteFunction
+    : public torch::autograd::Function<Qwen36FusedWeightedUnpermuteFunction> {
+    static int dtype_code(at::ScalarType type) {
+        switch (type) {
+            case at::kBFloat16: return 0;
+            case at::kFloat: return 1;
+            case at::kHalf: return 2;
+            default:
+                TORCH_CHECK(false,
+                    "fused MoE weighted unpermute does not support dtype ",
+                    type);
+        }
+    }
+
+    static at::Tensor forward(torch::autograd::AutogradContext* ctx,
+        at::Tensor returned, at::Tensor assignment_weights,
+        at::Tensor send_index, int64_t top_k) {
+        TORCH_CHECK(returned.is_cuda() && assignment_weights.is_cuda() &&
+                send_index.is_cuda(),
+            "fused MoE weighted unpermute expects CUDA tensors");
+        TORCH_CHECK(returned.device() == assignment_weights.device() &&
+                returned.device() == send_index.device(),
+            "fused MoE weighted unpermute tensors must share one CUDA device");
+        TORCH_CHECK(returned.dim() == 2 && assignment_weights.dim() == 1 &&
+                send_index.dim() == 1,
+            "invalid fused MoE weighted unpermute tensor ranks");
+        TORCH_CHECK(returned.is_contiguous() &&
+                assignment_weights.is_contiguous() && send_index.is_contiguous(),
+            "fused MoE weighted unpermute expects contiguous tensors");
+        TORCH_CHECK(send_index.scalar_type() == at::kLong,
+            "fused MoE weighted unpermute expects int64 send_index");
+        TORCH_CHECK(returned.scalar_type() == assignment_weights.scalar_type(),
+            "fused MoE weighted unpermute requires matching activation and "
+            "routing-weight dtypes");
+        TORCH_CHECK(returned.size(0) == assignment_weights.numel() &&
+                send_index.numel() == assignment_weights.numel(),
+            "fused MoE weighted unpermute assignment count mismatch");
+        TORCH_CHECK(top_k > 0 && assignment_weights.numel() % top_k == 0,
+            "invalid fused MoE weighted unpermute top_k");
+        TORCH_CHECK(top_k <= std::numeric_limits<int>::max() &&
+                assignment_weights.numel() <=
+                    std::numeric_limits<int>::max(),
+            "fused MoE weighted unpermute assignment grid is too large");
+
+        const int dtype = dtype_code(returned.scalar_type());
+        const int64_t assignments = assignment_weights.numel();
+        const int64_t tokens = assignments / top_k;
+        auto inverse_order = at::empty_like(send_index);
+        auto output = at::empty(
+            {tokens, returned.size(1)}, returned.options());
+        auto stream = c10::cuda::getCurrentCUDAStream(
+            returned.device().index()).stream();
+        launch_fused_moe_weighted_unpermute_forward(
+            returned.data_ptr(), assignment_weights.data_ptr(),
+            send_index.data_ptr<int64_t>(), inverse_order.data_ptr<int64_t>(),
+            output.data_ptr(), assignments, tokens, returned.size(1),
+            static_cast<int>(top_k), dtype, stream);
+        const auto launch_error = cudaGetLastError();
+        TORCH_CHECK(launch_error == cudaSuccess,
+            "fused MoE weighted unpermute forward launch failed: ",
+            cudaGetErrorString(launch_error));
+
+        ctx->save_for_backward(
+            {returned, assignment_weights, inverse_order});
+        ctx->saved_data["top_k"] = top_k;
+        ctx->saved_data["dtype"] = dtype;
+        return output;
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output) {
+        auto saved = ctx->get_saved_variables();
+        auto returned = saved[0];
+        auto assignment_weights = saved[1];
+        auto inverse_order = saved[2];
+        const int64_t top_k = ctx->saved_data["top_k"].toInt();
+        auto grad = grad_output[0].contiguous();
+        TORCH_CHECK(grad.scalar_type() == returned.scalar_type() &&
+                grad.dim() == 2 &&
+                grad.size(0) == assignment_weights.numel() / top_k &&
+                grad.size(1) == returned.size(1),
+            "invalid fused MoE weighted unpermute output gradient");
+        auto grad_returned = at::empty_like(returned);
+        auto grad_assignment_weights = at::empty_like(assignment_weights);
+        const int dtype = static_cast<int>(
+            ctx->saved_data["dtype"].toInt());
+        auto stream = c10::cuda::getCurrentCUDAStream(
+            returned.device().index()).stream();
+        launch_fused_moe_weighted_unpermute_backward(
+            returned.data_ptr(), assignment_weights.data_ptr(),
+            inverse_order.data_ptr<int64_t>(), grad.data_ptr(),
+            grad_returned.data_ptr(), grad_assignment_weights.data_ptr(),
+            assignment_weights.numel(), returned.size(1),
+            static_cast<int>(top_k), dtype, stream);
+        const auto launch_error = cudaGetLastError();
+        TORCH_CHECK(launch_error == cudaSuccess,
+            "fused MoE weighted unpermute backward launch failed: ",
+            cudaGetErrorString(launch_error));
+        return {grad_returned, grad_assignment_weights, at::Tensor(),
+            at::Tensor()};
+    }
+};
+
+struct Qwen36FusedLocalPermuteFunction
+    : public torch::autograd::Function<Qwen36FusedLocalPermuteFunction> {
+    static int dtype_code(at::ScalarType type) {
+        switch (type) {
+            case at::kBFloat16: return 0;
+            case at::kFloat: return 1;
+            case at::kHalf: return 2;
+            default:
+                TORCH_CHECK(false,
+                    "fused MoE local permutation does not support dtype ",
+                    type);
+        }
+    }
+
+    static std::vector<at::Tensor> forward(
+        torch::autograd::AutogradContext* ctx,
+        at::Tensor received, at::Tensor received_tokens,
+        at::Tensor received_experts, int64_t expert_count) {
+        TORCH_CHECK(received.is_cuda() && received_tokens.is_cuda() &&
+                received_experts.is_cuda(),
+            "fused MoE local permutation expects CUDA tensors");
+        TORCH_CHECK(received.device() == received_tokens.device() &&
+                received.device() == received_experts.device(),
+            "fused MoE local permutation tensors must share one CUDA device");
+        TORCH_CHECK(received.dim() == 2 && received_tokens.dim() == 1 &&
+                received_experts.dim() == 1,
+            "invalid fused MoE local permutation tensor ranks");
+        TORCH_CHECK(received.is_contiguous(),
+            "fused MoE local permutation expects contiguous activations");
+        TORCH_CHECK(received_tokens.scalar_type() == at::kLong &&
+                received_experts.scalar_type() == at::kLong,
+            "fused MoE local permutation expects int64 metadata");
+        TORCH_CHECK(received.size(0) == received_tokens.numel() &&
+                received.size(0) == received_experts.numel(),
+            "fused MoE local permutation row count mismatch");
+        TORCH_CHECK(expert_count > 0 && expert_count <= 1024,
+            "fused MoE local permutation supports 1..1024 local experts");
+        TORCH_CHECK(received.size(0) <= std::numeric_limits<int>::max(),
+            "fused MoE local permutation has too many received rows");
+
+        const int dtype = dtype_code(received.scalar_type());
+        auto selected = at::empty_like(received);
+        auto selected_tokens = at::empty(
+            {received.size(0)}, received_tokens.options());
+        auto sorted_experts = at::empty(
+            {received.size(0)}, received_experts.options());
+        auto int_options = at::TensorOptions()
+            .device(received.device()).dtype(at::kInt);
+        auto counts = at::empty({expert_count}, int_options);
+        auto offsets = at::empty({expert_count}, int_options);
+        auto cursors = at::empty({expert_count}, int_options);
+        auto sorted_to_received = at::empty(
+            {received.size(0)}, received_experts.options());
+        auto stream = c10::cuda::getCurrentCUDAStream(
+            received.device().index()).stream();
+        launch_fused_moe_local_permute_forward(
+            received.data_ptr(), received_tokens.data_ptr<int64_t>(),
+            received_tokens.stride(0), received_experts.data_ptr<int64_t>(),
+            received_experts.stride(0), selected.data_ptr(),
+            selected_tokens.data_ptr<int64_t>(),
+            sorted_experts.data_ptr<int64_t>(), counts.data_ptr<int>(),
+            offsets.data_ptr<int>(), cursors.data_ptr<int>(),
+            sorted_to_received.data_ptr<int64_t>(), received.size(0),
+            received.size(1), static_cast<int>(expert_count), dtype, stream);
+        const auto launch_error = cudaGetLastError();
+        TORCH_CHECK(launch_error == cudaSuccess,
+            "fused MoE local permutation launch failed: ",
+            cudaGetErrorString(launch_error));
+
+        ctx->save_for_backward({received, sorted_to_received});
+        ctx->saved_data["dtype"] = dtype;
+        ctx->saved_data["hidden"] = received.size(1);
+        ctx->mark_non_differentiable(
+            {selected_tokens, sorted_experts, counts, offsets,
+                sorted_to_received});
+        return {selected, selected_tokens, sorted_experts, counts, offsets,
+            sorted_to_received};
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output) {
+        auto saved = ctx->get_saved_variables();
+        const auto& received = saved[0];
+        const auto& sorted_to_received = saved[1];
+        if (!grad_output[0].defined()) {
+            return {at::zeros_like(received), at::Tensor(), at::Tensor(),
+                at::Tensor()};
+        }
+        auto grad_selected = grad_output[0].contiguous();
+        auto grad_received = at::empty_like(grad_selected);
+        const int dtype = static_cast<int>(
+            ctx->saved_data["dtype"].toInt());
+        const int64_t hidden = ctx->saved_data["hidden"].toInt();
+        auto stream = c10::cuda::getCurrentCUDAStream(
+            grad_selected.device().index()).stream();
+        launch_fused_moe_local_permute_rows(
+            grad_selected.data_ptr(), sorted_to_received.data_ptr<int64_t>(),
+            grad_received.data_ptr(), sorted_to_received.numel(), hidden,
+            dtype, true, stream);
+        const auto launch_error = cudaGetLastError();
+        TORCH_CHECK(launch_error == cudaSuccess,
+            "fused MoE local permutation backward launch failed: ",
+            cudaGetErrorString(launch_error));
+        return {grad_received, at::Tensor(), at::Tensor(), at::Tensor()};
+    }
+};
+
+struct Qwen36FusedLocalInversePermuteFunction
+    : public torch::autograd::Function<
+          Qwen36FusedLocalInversePermuteFunction> {
+    static at::Tensor forward(torch::autograd::AutogradContext* ctx,
+        at::Tensor sorted_output, at::Tensor sorted_to_received,
+        int64_t dtype) {
+        TORCH_CHECK(sorted_output.is_cuda() &&
+                sorted_to_received.is_cuda() && sorted_output.is_contiguous(),
+            "fused MoE local inverse permutation expects contiguous CUDA "
+            "tensors");
+        TORCH_CHECK(sorted_output.device() == sorted_to_received.device() &&
+                sorted_to_received.is_contiguous(),
+            "fused MoE local inverse permutation tensors must be contiguous "
+            "and share one CUDA device");
+        TORCH_CHECK(sorted_output.dim() == 2 &&
+                sorted_to_received.dim() == 1 &&
+                sorted_to_received.scalar_type() == at::kLong &&
+                sorted_output.size(0) == sorted_to_received.numel(),
+            "invalid fused MoE local inverse permutation tensors");
+        TORCH_CHECK(dtype >= 0 && dtype <= 2 &&
+                dtype == Qwen36FusedLocalPermuteFunction::dtype_code(
+                    sorted_output.scalar_type()),
+            "invalid fused MoE local inverse permutation dtype code");
+        auto output = at::empty_like(sorted_output);
+        auto stream = c10::cuda::getCurrentCUDAStream(
+            sorted_output.device().index()).stream();
+        launch_fused_moe_local_permute_rows(
+            sorted_output.data_ptr(),
+            sorted_to_received.data_ptr<int64_t>(), output.data_ptr(),
+            sorted_output.size(0), sorted_output.size(1),
+            static_cast<int>(dtype), true, stream);
+        const auto launch_error = cudaGetLastError();
+        TORCH_CHECK(launch_error == cudaSuccess,
+            "fused MoE local inverse permutation launch failed: ",
+            cudaGetErrorString(launch_error));
+        ctx->save_for_backward({sorted_output, sorted_to_received});
+        ctx->saved_data["dtype"] = dtype;
+        ctx->saved_data["hidden"] = sorted_output.size(1);
+        return output;
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output) {
+        auto saved = ctx->get_saved_variables();
+        const auto& sorted_output = saved[0];
+        const auto& sorted_to_received = saved[1];
+        if (!grad_output[0].defined()) {
+            return {at::zeros_like(sorted_output), at::Tensor(),
+                at::Tensor()};
+        }
+        auto grad_received_order = grad_output[0].contiguous();
+        auto grad_sorted = at::empty_like(grad_received_order);
+        const int dtype = static_cast<int>(
+            ctx->saved_data["dtype"].toInt());
+        const int64_t hidden = ctx->saved_data["hidden"].toInt();
+        auto stream = c10::cuda::getCurrentCUDAStream(
+            grad_received_order.device().index()).stream();
+        launch_fused_moe_local_permute_rows(
+            grad_received_order.data_ptr(),
+            sorted_to_received.data_ptr<int64_t>(), grad_sorted.data_ptr(),
+            sorted_to_received.numel(), hidden, dtype, false, stream);
+        const auto launch_error = cudaGetLastError();
+        TORCH_CHECK(launch_error == cudaSuccess,
+            "fused MoE local inverse permutation backward launch failed: ",
+            cudaGetErrorString(launch_error));
+        return {grad_sorted, at::Tensor(), at::Tensor()};
+    }
+};
+
+struct RoutedExpertLora {
+    const at::Tensor* gate_up_a = nullptr;  // [local_experts, rank, hidden]
+    const at::Tensor* gate_up_b = nullptr;  // [local_experts, 2*intermediate, rank]
+    const at::Tensor* down_a = nullptr;     // [local_experts, rank, intermediate]
+    const at::Tensor* down_b = nullptr;     // [local_experts, hidden, rank]
+    double scaling = 0.0;
+};
+
+enum class LoraTpLayout : uint8_t {
+    LatentRank,
+    ColumnParallel,
+    RowParallel,
+};
+
+struct TrainingContext;
+
+static LoraTpLayout lora_tp_layout(
+    const TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx);
+
+// Per-sample adapter projection used by the batched multi-LoRA path.  This is
+// intentionally separate from RoutedExpertLora: routed experts carry one
+// A/B pair per local expert, while dense/shared projections carry one pair per
+// adapter sample.
+struct LoraBatchEntry {
+    at::Tensor a_stack;    // [N, rank, in]
+    at::Tensor b_stack;    // [N, out, rank]
+    at::Tensor scaling;    // [N, 1, 1]
+    LoraTpLayout layout = LoraTpLayout::LatentRank;
+};
+
+static const LoraBatchEntry* lora_batch_entry(
+    TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx);
+static at::Tensor dense_mlp_forward_batched(
+    TrainingContext* ctx, int64_t layer_idx, const at::Tensor& hidden,
+    const at::Tensor& gate_proj, const at::Tensor& up_proj,
+    const at::Tensor& down_proj, at::ScalarType compute_type,
+    const LayerConfig* config_override = nullptr);
+static at::Tensor fused_mlp_fc1_weight(
+    TrainingContext* ctx, int64_t layer_idx);
+static bool base_tp_mlp_enabled(const TrainingContext* ctx);
+static int64_t base_tp_mlp_world_size(const TrainingContext* ctx);
+static at::Tensor tp_allreduce_base_mlp(
+    TrainingContext* ctx, const at::Tensor& local_output);
+static at::Tensor tp_copy_base_mlp_input(
+    TrainingContext* ctx, const at::Tensor& input);
+
+static at::Tensor lora_activation_delta(
+    TrainingContext* ctx, const at::Tensor& x,
+    const at::Tensor& A, const at::Tensor& B,
+    const at::Tensor& scaling, LoraTpLayout layout);
+
+static at::Tensor add_batched_lora(
+    TrainingContext* ctx, const at::Tensor& base, const at::Tensor& input,
+    const LoraBatchEntry* entry
+) {
+    if (!entry) return base;
+    return base + lora_activation_delta(
+        ctx, input, entry->a_stack, entry->b_stack, entry->scaling,
+        entry->layout);
+}
+
+// Per-token routed-expert LoRA. Dynamic adapters add a leading sample axis to
+// the expert-local tensors: A [batch, experts, rank, in],
+// B [batch, experts, out, rank]. Flattening (sample, expert) lets one pair of
+// index_select + bmm operations select the correct adapter and expert without
+// materializing a full-rank delta weight.
+static at::Tensor dynamic_expert_lora_delta(
+    TrainingContext* ctx,
+    const at::Tensor& input,
+    const at::Tensor& token_indices,
+    const at::Tensor& local_expert_indices,
+    int64_t batch,
+    int64_t seq,
+    const LoraBatchEntry* entry
+) {
+    if (!entry) return at::zeros({0}, input.options());
+    TORCH_CHECK(entry->a_stack.dim() == 4 && entry->b_stack.dim() == 4,
+        "dynamic routed-expert LoRA expects rank-4 stacked A/B tensors");
+    const int64_t local_experts = entry->a_stack.size(1);
+    auto sample_indices = at::floor_divide(token_indices, seq);
+    auto pair_indices = sample_indices * local_experts + local_expert_indices;
+    auto a_stack = entry->a_stack;
+    auto b_stack = entry->b_stack;
+    if (a_stack.size(0) == 1 && batch > 1) {
+        a_stack = a_stack.expand(
+            {batch, a_stack.size(1), a_stack.size(2), a_stack.size(3)});
+        b_stack = b_stack.expand(
+            {batch, b_stack.size(1), b_stack.size(2), b_stack.size(3)});
+    }
+    auto a = a_stack.flatten(0, 1)
+        .index_select(0, pair_indices).to(input.scalar_type());
+    auto b = b_stack.flatten(0, 1)
+        .index_select(0, pair_indices).to(input.scalar_type());
+    auto lora_input = entry->layout == LoraTpLayout::LatentRank
+        ? tp_copy_lora_input(ctx, input) : input;
+    auto low_rank = at::bmm(a, lora_input.unsqueeze(-1)).squeeze(-1);
+    auto delta = at::bmm(b, low_rank.unsqueeze(-1)).squeeze(-1);
+    auto scaling_stack = entry->scaling;
+    if (scaling_stack.size(0) == 1 && batch > 1) {
+        scaling_stack = scaling_stack.expand({batch, 1, 1});
+    }
+    auto scaling = scaling_stack.index_select(0, sample_indices)
+        .reshape({-1, 1}).to(input.scalar_type());
+    auto scaled = delta * scaling;
+    return entry->layout == LoraTpLayout::LatentRank
+        ? tp_allreduce_lora_delta(ctx, scaled) : scaled;
+}
+
+static at::Tensor moe_routed_a2a(
+    TrainingContext* training_ctx, ncclComm_t comm,
+    const at::Tensor& flat, const at::Tensor& topk_weights,
+    const at::Tensor& topk_indices,
+    const at::Tensor& experts_gate_up, const at::Tensor& experts_down,
+    const RoutedExpertLora& expert_lora,
+    int64_t top_k, int64_t intermediate, int64_t expert_count,
+    int64_t batch, int64_t seq,
+    cudaStream_t requested_stream,
+    const LoraBatchEntry* expert_gate_up_lora,
+    const LoraBatchEntry* expert_down_lora
+) {
+    const int64_t hidden_dim = flat.size(1);
+    const auto metadata_stream = env_enabled(
+            "QWEN36_EP_A2A_COUNT_OVERLAP")
+        ? context_moe_ep_metadata_stream(training_ctx)
+        : nullptr;
+    const bool expert_tp = base_tp_mlp_enabled(training_ctx);
+    const auto gate_up_layout = expert_tp
+        ? LoraTpLayout::ColumnParallel : LoraTpLayout::LatentRank;
+    const auto down_layout = expert_tp
+        ? LoraTpLayout::RowParallel : LoraTpLayout::LatentRank;
+    auto routed_output = at::zeros_like(flat);
+    auto token_indices = at::arange(
+        flat.size(0), at::TensorOptions().device(flat.device()).dtype(at::kLong));
+    auto run_local_experts = [&](const at::Tensor& received,
+                                 const at::Tensor& received_tokens,
+                                 const at::Tensor& received_experts) {
+        at::Tensor sorted_experts;
+        at::Tensor expert_order;
+        at::Tensor selected;
+        at::Tensor selected_tokens;
+        at::Tensor counts;
+        at::Tensor offsets;
+        const bool fused_local_permute =
+            env_enabled("QWEN36_MOE_FUSED_LOCAL_PERMUTE") &&
+            expert_count <= 1024;
+        if (fused_local_permute) {
+            auto permutation = Qwen36FusedLocalPermuteFunction::apply(
+                received, received_tokens, received_experts, expert_count);
+            selected = permutation[0];
+            selected_tokens = permutation[1];
+            sorted_experts = permutation[2];
+            counts = permutation[3];
+            offsets = permutation[4];
+            expert_order = permutation[5];
+        } else {
+            auto sorted = received_experts.sort(0);
+            sorted_experts = std::get<0>(sorted);
+            expert_order = std::get<1>(sorted);
+            selected = received.index_select(0, expert_order);
+            selected_tokens = received_tokens.index_select(0, expert_order);
+            counts = at::bincount(
+                sorted_experts, c10::nullopt, expert_count);
+            offsets = counts.cumsum(0).to(at::kInt);
+        }
+        auto sorted_output = at::zeros_like(selected);
+
+        if (selected.size(0) > 0) {
+            at::Tensor expert_out;
+#if RUSTRAIN_HAS_ATEN_GROUPED_MM
+            const bool use_grouped_mm =
+                !env_enabled("QWEN36_DISABLE_GROUPED_MM") &&
+                selected.scalar_type() == at::kBFloat16 &&
+                hidden_dim % 8 == 0 && intermediate % 8 == 0;
+            if (use_grouped_mm) {
+                at::Tensor counts_cpu;
+                auto fixed_lora_delta = [&](const at::Tensor& input,
+                                            const at::Tensor& a,
+                                            const at::Tensor& b,
+                                            LoraTpLayout layout) {
+                    auto lora_input = layout == LoraTpLayout::LatentRank
+                        ? tp_copy_lora_input(training_ctx, input) : input;
+                    at::Tensor delta;
+                    if (a.size(1) % 8 == 0) {
+                        auto low_rank = at::_grouped_mm(
+                            lora_input, a.transpose(1, 2), offsets);
+                        delta = at::_grouped_mm(
+                            low_rank, b.transpose(1, 2), offsets);
+                    } else {
+                        if (!counts_cpu.defined()) {
+                            counts_cpu = counts.to(
+                                at::TensorOptions().device(at::kCPU));
+                        }
+                        std::vector<at::Tensor> chunks;
+                        chunks.reserve(expert_count);
+                        int64_t offset = 0;
+                        for (int64_t e_local = 0;
+                             e_local < expert_count; ++e_local) {
+                            const int64_t rows = counts_cpu.index(
+                                {e_local}).item<int64_t>();
+                            auto expert_input = lora_input.narrow(
+                                0, offset, rows);
+                            chunks.push_back(at::matmul(
+                                at::matmul(
+                                    expert_input, a.select(0, e_local).t()),
+                                b.select(0, e_local).t()));
+                            offset += rows;
+                        }
+                        delta = at::cat(chunks, 0);
+                    }
+                    auto scaled = delta * expert_lora.scaling;
+                    return layout == LoraTpLayout::LatentRank
+                        ? tp_allreduce_lora_delta(training_ctx, scaled)
+                        : scaled;
+                };
+                auto gu = at::_grouped_mm(
+                    selected, experts_gate_up.transpose(1, 2), offsets);
+                if (expert_lora.gate_up_a && expert_lora.gate_up_b) {
+                    gu = gu + fixed_lora_delta(
+                        selected, *expert_lora.gate_up_a,
+                        *expert_lora.gate_up_b, gate_up_layout);
+                }
+                if (expert_gate_up_lora) {
+                    gu = gu + dynamic_expert_lora_delta(
+                        training_ctx, selected, selected_tokens,
+                        sorted_experts, batch, seq, expert_gate_up_lora);
+                }
+                auto activated = fused_swiglu_op(
+                    gu.narrow(-1, 0, intermediate),
+                    gu.narrow(-1, intermediate, intermediate), 0.0);
+                expert_out = at::_grouped_mm(
+                    activated, experts_down.transpose(1, 2), offsets);
+                if (expert_lora.down_a && expert_lora.down_b) {
+                    expert_out = expert_out + fixed_lora_delta(
+                        activated, *expert_lora.down_a,
+                        *expert_lora.down_b, down_layout);
+                }
+                if (expert_down_lora) {
+                    expert_out = expert_out + dynamic_expert_lora_delta(
+                        training_ctx, activated, selected_tokens,
+                        sorted_experts, batch, seq, expert_down_lora);
+                }
+            } else
+#endif
+            {
+                auto counts_cpu = counts.to(
+                    at::TensorOptions().device(at::kCPU));
+                expert_out = at::zeros_like(selected);
+                int64_t offset = 0;
+                for (int64_t e_local = 0; e_local < expert_count; ++e_local) {
+                    const int64_t rows =
+                        counts_cpu.index({e_local}).item<int64_t>();
+                    if (rows > 0) {
+                        auto expert_input = selected.narrow(0, offset, rows);
+                        auto expert_tokens = selected_tokens.narrow(
+                            0, offset, rows);
+                        auto local_experts = sorted_experts.narrow(
+                            0, offset, rows);
+                        auto gu = at::matmul(
+                            expert_input,
+                            experts_gate_up.select(0, e_local).t());
+                        if (expert_lora.gate_up_a && expert_lora.gate_up_b) {
+                            auto lora_input = gate_up_layout ==
+                                    LoraTpLayout::LatentRank
+                                ? tp_copy_lora_input(training_ctx, expert_input)
+                                : expert_input;
+                            auto delta = at::matmul(
+                                at::matmul(lora_input,
+                                    expert_lora.gate_up_a->select(
+                                        0, e_local).t()),
+                                expert_lora.gate_up_b->select(0, e_local).t()) *
+                                expert_lora.scaling;
+                            gu = gu + (gate_up_layout ==
+                                    LoraTpLayout::LatentRank
+                                ? tp_allreduce_lora_delta(training_ctx, delta)
+                                : delta);
+                        }
+                        if (expert_gate_up_lora) {
+                            gu = gu + dynamic_expert_lora_delta(
+                                training_ctx, expert_input, expert_tokens,
+                                local_experts, batch, seq,
+                                expert_gate_up_lora);
+                        }
+                        auto activated = fused_swiglu_op(
+                            gu.narrow(-1, 0, intermediate),
+                            gu.narrow(-1, intermediate, intermediate), 0.0);
+                        auto local_out = at::matmul(
+                            activated,
+                            experts_down.select(0, e_local).t());
+                        if (expert_lora.down_a && expert_lora.down_b) {
+                            auto lora_input = down_layout ==
+                                    LoraTpLayout::LatentRank
+                                ? tp_copy_lora_input(training_ctx, activated)
+                                : activated;
+                            auto delta = at::matmul(
+                                at::matmul(lora_input,
+                                    expert_lora.down_a->select(
+                                        0, e_local).t()),
+                                expert_lora.down_b->select(0, e_local).t()) *
+                                expert_lora.scaling;
+                            local_out = local_out + (down_layout ==
+                                    LoraTpLayout::LatentRank
+                                ? tp_allreduce_lora_delta(training_ctx, delta)
+                                : delta);
+                        }
+                        if (expert_down_lora) {
+                            local_out = local_out + dynamic_expert_lora_delta(
+                                training_ctx, activated, expert_tokens,
+                                local_experts, batch, seq, expert_down_lora);
+                        }
+                        expert_out = expert_out.index_add(
+                            0,
+                            at::arange(offset, offset + rows,
+                                expert_order.options()),
+                            local_out);
+                    }
+                    offset += rows;
+                }
+            }
+            sorted_output = expert_out;
+        }
+
+        auto local_output = fused_local_permute
+            ? Qwen36FusedLocalInversePermuteFunction::apply(
+                sorted_output, expert_order,
+                Qwen36FusedLocalPermuteFunction::dtype_code(
+                    sorted_output.scalar_type()))
+            : at::zeros_like(received).index_add(
+                0, expert_order, sorted_output);
+        // Empty destinations still need activation and parameter graph edges
+        // so every rank reaches the same optimizer collectives.
+        if (!local_output.requires_grad()) {
+            at::Tensor anchor = received.sum().to(local_output.scalar_type());
+            auto include = [&](const at::Tensor* tensor) {
+                if (!tensor || !tensor->defined() ||
+                    !tensor->requires_grad()) return;
+                anchor = anchor + tensor->sum().to(local_output.scalar_type());
+            };
+            include(expert_lora.gate_up_a);
+            include(expert_lora.gate_up_b);
+            include(expert_lora.down_a);
+            include(expert_lora.down_b);
+            if (expert_gate_up_lora) {
+                include(&expert_gate_up_lora->a_stack);
+                include(&expert_gate_up_lora->b_stack);
+            }
+            if (expert_down_lora) {
+                include(&expert_down_lora->a_stack);
+                include(&expert_down_lora->b_stack);
+            }
+            local_output = local_output + anchor * 0.0;
+        }
+        return local_output;
+    };
+
+    auto dispatch_and_combine = [&](const at::Tensor& assignment_input,
+                                    const at::Tensor& assignment_experts,
+                                    const at::Tensor& assignment_tokens,
+                                    const at::Tensor& assignment_weights,
+                                    bool packed_assignments) {
+        // `received_tokens` preserves the source flattened row index through
+        // dispatch so dynamic multi-LoRA can recover the tenant/sample row.
+        auto dispatched = Qwen36A2ADispatchFunction::apply(
+            assignment_input, assignment_experts, assignment_tokens,
+            expert_count,
+            static_cast<int64_t>(reinterpret_cast<uintptr_t>(comm)),
+            static_cast<int64_t>(reinterpret_cast<uintptr_t>(requested_stream)),
+            static_cast<int64_t>(reinterpret_cast<uintptr_t>(metadata_stream)));
+        auto received = dispatched[0];
+        auto received_tokens = dispatched[1];
+        auto received_experts = dispatched[2];
+        auto send_index = dispatched[3];
+        auto send_offsets = dispatched[4];
+        auto recv_offsets = dispatched[5];
+        auto local_output = run_local_experts(
+            received, received_tokens, received_experts);
+        auto returned = Qwen36A2ACombineFunction::apply(
+            local_output, send_offsets, recv_offsets,
+            static_cast<int64_t>(reinterpret_cast<uintptr_t>(comm)),
+            static_cast<int64_t>(reinterpret_cast<uintptr_t>(requested_stream)));
+        returned = tp_allreduce_base_mlp(training_ctx, returned);
+        if (packed_assignments &&
+            env_enabled("QWEN36_MOE_FUSED_UNPERMUTE")) {
+            routed_output = Qwen36FusedWeightedUnpermuteFunction::apply(
+                returned, assignment_weights, send_index, top_k);
+        } else {
+            auto source_tokens = assignment_tokens.index_select(0, send_index);
+            auto source_weights = assignment_weights
+                .index_select(0, send_index).unsqueeze(-1);
+            routed_output = routed_output.index_add(
+                0, source_tokens, returned * source_weights);
+        }
+    };
+
+    if (env_enabled("QWEN36_EP_A2A_PACKED", true)) {
+        auto assignment_tokens = token_indices.unsqueeze(1)
+            .expand({flat.size(0), top_k}).reshape({-1});
+        auto assignment_experts = topk_indices.reshape({-1}).contiguous();
+        auto assignment_weights = topk_weights.reshape({-1}).contiguous();
+        dispatch_and_combine(
+            flat, assignment_experts, assignment_tokens,
+            assignment_weights, true);
+    } else {
+        for (int64_t kk = 0; kk < top_k; ++kk) {
+            auto expert_indices = topk_indices.select(-1, kk).contiguous();
+            auto expert_weights = topk_weights.select(-1, kk).contiguous();
+            dispatch_and_combine(
+                flat, expert_indices, token_indices, expert_weights, false);
+        }
+    }
+    return routed_output;
+}
+
+static at::Tensor moe_forward(
+    TrainingContext* training_ctx,
+    int64_t layer_idx,
+    void* nccl_comm_v, void* nccl_stream_v,
+    const at::Tensor& hidden,
+    const at::Tensor& gate_w, const at::Tensor& shared_expert_gate_w,
+    const at::Tensor& shared_gate_proj, const at::Tensor& shared_up_proj, const at::Tensor& shared_down_proj,
+    const at::Tensor& experts_gate_up, const at::Tensor& experts_down,
+    const RoutedExpertLora& expert_lora,
+    int64_t num_experts, int64_t top_k, int64_t intermediate,
+    bool norm_topk_prob, int64_t expert_start, int64_t expert_count,
+    at::ScalarType compute_type, bool use_batched,
+    const LoraBatchEntry* shared_gate_lora = nullptr,
+    const LoraBatchEntry* shared_up_lora = nullptr,
+    const LoraBatchEntry* shared_down_lora = nullptr,
+    const LoraBatchEntry* expert_gate_up_lora = nullptr,
+    const LoraBatchEntry* expert_down_lora = nullptr
+) {
+    int64_t batch = hidden.size(0), seq = hidden.size(1), hidden_dim = hidden.size(2);
+    auto device = hidden.device();
+    auto flat = hidden.reshape({batch * seq, hidden_dim});
+    int64_t N = flat.size(0);
+    const bool expert_tp = base_tp_mlp_enabled(training_ctx);
+    const int64_t local_intermediate = experts_gate_up.size(1) / 2;
+    TORCH_CHECK(experts_gate_up.dim() == 3 && experts_down.dim() == 3 &&
+            experts_gate_up.size(0) == expert_count &&
+            experts_down.size(0) == expert_count &&
+            experts_gate_up.size(1) == 2 * local_intermediate &&
+            experts_gate_up.size(2) == hidden_dim &&
+            experts_down.size(1) == hidden_dim &&
+            experts_down.size(2) == local_intermediate,
+        "routed expert local weight shapes are inconsistent: gate_up=",
+        experts_gate_up.sizes(), " down=", experts_down.sizes());
+    if (expert_tp) {
+        const int64_t tp_world_size = base_tp_mlp_world_size(training_ctx);
+        TORCH_CHECK(intermediate > 0 &&
+                intermediate % tp_world_size == 0 &&
+                local_intermediate == intermediate / tp_world_size,
+            "routed expert TP intermediate mismatch: global=", intermediate,
+            " local=", local_intermediate,
+            " TP_SIZE=", tp_world_size);
+    } else {
+        TORCH_CHECK(local_intermediate == intermediate,
+            "unsharded routed expert intermediate mismatch: configured=",
+            intermediate, " weight=", local_intermediate);
+    }
+    auto expert_flat = tp_copy_base_mlp_input(training_ctx, flat);
+
+    auto compute_shared_expert = [&]() {
+        auto shared_input = tp_copy_base_mlp_input(training_ctx, flat);
+        at::Tensor shared_gate;
+        at::Tensor shared_up;
+        const auto fused_shared_fc1 = use_batched
+            ? fused_mlp_fc1_weight(training_ctx, layer_idx)
+            : at::Tensor();
+        if (fused_shared_fc1.defined()) {
+            TORCH_CHECK(fused_shared_fc1.dim() == 2 &&
+                    fused_shared_fc1.size(0) ==
+                        shared_gate_proj.size(0) + shared_up_proj.size(0) &&
+                    fused_shared_fc1.size(1) == shared_input.size(1),
+                "fused shared-expert FC1 weight shape is incompatible with the "
+                "local MLP geometry");
+            auto shared_fc1 = at::matmul(
+                shared_input, fused_shared_fc1.t());
+            shared_gate = shared_fc1.narrow(
+                -1, 0, shared_gate_proj.size(0));
+            shared_up = shared_fc1.narrow(
+                -1, shared_gate_proj.size(0), shared_up_proj.size(0));
+        } else {
+            shared_gate = at::matmul(shared_input, shared_gate_proj.t());
+            shared_up = at::matmul(shared_input, shared_up_proj.t());
+        }
+        if (shared_gate_lora) {
+            shared_gate = add_batched_lora(
+                training_ctx, shared_gate.reshape({batch, seq, -1}),
+                shared_input.reshape({batch, seq, hidden_dim}),
+                shared_gate_lora)
+                .reshape({batch * seq, -1});
+        }
+        if (shared_up_lora) {
+            shared_up = add_batched_lora(
+                training_ctx, shared_up.reshape({batch, seq, -1}),
+                shared_input.reshape({batch, seq, hidden_dim}),
+                shared_up_lora)
+                .reshape({batch * seq, -1});
+        }
+        auto shared_hidden = fused_swiglu_op(
+            shared_gate.reshape({batch, seq, -1}),
+            shared_up.reshape({batch, seq, -1}), 0.0);
+        auto shared_out = at::matmul(
+            shared_hidden.reshape({batch * seq, -1}),
+            shared_down_proj.t());
+        if (shared_down_lora) {
+            shared_out = add_batched_lora(
+                training_ctx, shared_out.reshape({batch, seq, -1}),
+                shared_hidden, shared_down_lora)
+                .reshape({batch * seq, -1});
+        }
+        shared_out = tp_allreduce_base_mlp(training_ctx, shared_out);
+        auto seg = at::sigmoid(
+            at::matmul(flat, shared_expert_gate_w.t())).to(compute_type);
+        return (shared_out * seg).to(compute_type);
+    };
+
+    at::Tensor shared_out;
+    Qwen36EventPair shared_events;
+    bool shared_overlap_active = false;
+    if (env_enabled("QWEN36_MOE_SHARED_OVERLAP")) {
+        // The shared expert consumes only the replicated flattened input. Its
+        // dense GEMMs can therefore run while the current stream performs
+        // routed sorting/A2A and local expert grouped GEMMs. The event pair
+        // avoids a device-wide synchronize and is kept opt-in until each
+        // target topology has a matched latency result.
+        const auto current_stream =
+            c10::cuda::getCurrentCUDAStream(device.index()).stream();
+        const auto raw_shared_stream =
+            context_moe_shared_stream(training_ctx);
+        const auto shared_stream = c10::cuda::getStreamFromExternal(
+            raw_shared_stream, device.index());
+        TORCH_CHECK(cudaEventCreateWithFlags(
+                &shared_events.ready, cudaEventDisableTiming) == cudaSuccess,
+            "failed to create shared-expert producer event");
+        TORCH_CHECK(cudaEventCreateWithFlags(
+                &shared_events.done, cudaEventDisableTiming) == cudaSuccess,
+            "failed to create shared-expert completion event");
+        TORCH_CHECK(cudaEventRecord(shared_events.ready, current_stream) ==
+                cudaSuccess,
+            "failed to record shared-expert producer event");
+        TORCH_CHECK(cudaStreamWaitEvent(
+                shared_stream.stream(), shared_events.ready, 0) == cudaSuccess,
+            "failed to fence shared-expert stream");
+        {
+            c10::cuda::CUDAStreamGuard guard(shared_stream);
+            shared_out = compute_shared_expert();
+        }
+        TORCH_CHECK(cudaEventRecord(
+                shared_events.done, shared_stream.stream()) == cudaSuccess,
+            "failed to record shared-expert completion event");
+        shared_overlap_active = true;
+    }
+
+    const bool sharded_a2a_mode =
+        env_enabled("QWEN36_EP_A2A_SHARDED") &&
+        expert_count < num_experts;
+
+    auto router_logits = at::matmul(flat, gate_w.t());
+    auto routing_weights = router_logits.softmax(-1, at::kFloat);
+    auto [topk_weights, topk_indices] = routing_weights.topk(top_k, -1, true, true);
+    routing_weights = attach_router_aux_loss(
+        training_ctx, routing_weights, topk_indices,
+        batch, seq, top_k, num_experts);
+    topk_weights = routing_weights.gather(-1, topk_indices);
+    if (norm_topk_prob) {
+        auto denom = topk_weights.sum(-1, true).clamp_min(1e-9);
+        topk_weights = topk_weights / denom;
+    }
+    topk_weights = topk_weights.to(compute_type);
+
+    auto routed_output = at::zeros(flat.sizes(), flat.options());
+    if (env_enabled("QWEN36_EP_A2A_SHARDED") && expert_count < num_experts) {
+        TORCH_CHECK(nccl_comm_v && env_enabled("QWEN36_EP_A2A"),
+            "QWEN36_EP_A2A_SHARDED=1 requires an initialized EP communicator "
+            "and QWEN36_EP_A2A=1");
+    }
+    bool use_a2a = false;
+    if (expert_tp && expert_count < num_experts) {
+        TORCH_CHECK(sharded_a2a_mode && env_enabled("QWEN36_EP_A2A") &&
+                env_enabled("QWEN36_EP_A2A_PACKED", true),
+            "expert TP with expert sharding requires packed sharded A2A: set "
+            "QWEN36_EP_A2A=1, QWEN36_EP_A2A_SHARDED=1, and keep "
+            "QWEN36_EP_A2A_PACKED enabled");
+    }
+    if (nccl_comm_v && env_enabled("QWEN36_EP_A2A") &&
+        ((!expert_gate_up_lora && !expert_down_lora) || sharded_a2a_mode)) {
+        auto comm = reinterpret_cast<ncclComm_t>(nccl_comm_v);
+        int world = 1, rank = 0;
+        auto err = ncclCommCount(comm, &world);
+        TORCH_CHECK(err == ncclSuccess, "ncclCommCount failed: ",
+            ncclGetErrorString(err));
+        err = ncclCommUserRank(comm, &rank);
+        TORCH_CHECK(err == ncclSuccess, "ncclCommUserRank failed: ",
+            ncclGetErrorString(err));
+        TORCH_CHECK(world * expert_count == num_experts,
+            "EP A2A requires equal contiguous expert partitions: world=", world,
+            " local_experts=", expert_count, " global_experts=", num_experts);
+        TORCH_CHECK(expert_start == rank * expert_count,
+            "EP A2A requires rank-contiguous expert ownership: rank=", rank,
+            " expert_start=", expert_start, " local_experts=", expert_count);
+        use_a2a = world > 1;
+        if (use_a2a) {
+            routed_output = moe_routed_a2a(
+                training_ctx, comm, expert_flat, topk_weights, topk_indices,
+                experts_gate_up, experts_down, expert_lora,
+                top_k, local_intermediate, expert_count, batch, seq,
+                reinterpret_cast<cudaStream_t>(nccl_stream_v),
+                expert_gate_up_lora, expert_down_lora);
+        }
+    }
+
+    // Debug: dump MoE routing and weight stats
+    if (getenv("QWEN36_DUMP_MOE")) {
+        auto rl_f = router_logits.to(at::kFloat);
+        auto rw_f = topk_weights.to(at::kFloat);
+        auto egu_f = experts_gate_up.select(0, 0).to(at::kFloat);
+        auto ed_f = experts_down.select(0, 0).to(at::kFloat);
+        auto sg_f = shared_gate_proj.to(at::kFloat);
+        auto sd_f = shared_down_proj.to(at::kFloat);
+        auto seg_f = at::sigmoid(at::matmul(flat, shared_expert_gate_w.t())).to(at::kFloat);
+    }
+
+    // Sort-based expert dispatch — eliminates eq/nonzero/index_select per expert.
+    // At large N (1000+), the old for-loop with eq+nonzero per expert was O(experts × tokens)
+    // with high kernel launch overhead. This approach pre-sorts tokens by expert assignment.
+    //
+    // Autograd note: routing indices/weights are computed in no-grad forward (detached).
+    // Only matmul inputs/outputs participate in autograd. sort/index_select/index_add
     // gradients are handled by PyTorch automatically.
-    for (int64_t kk = 0; kk < top_k; kk++) {
+    if (!use_a2a) for (int64_t kk = 0; kk < top_k; kk++) {
         auto expert_indices = topk_indices.select(-1, kk);   // [N*seq]
         auto expert_weights = topk_weights.select(-1, kk);   // [N*seq]
 
@@ -697,44 +3231,228 @@ static at::Tensor moe_forward(
         // Find expert boundaries via bincount + cumsum
         auto counts = at::bincount(sorted_indices, c10::nullopt, expert_start + expert_count);
         // counts[e_global] = number of tokens assigned to expert e
-
         // Gather tokens in sorted order (contiguous per expert)
-        auto gathered = flat.index_select(0, sort_order);
-        auto gathered_weights = expert_weights.index_select(0, sort_order).unsqueeze(-1);
-
-        // Process each expert's contiguous token slice
-        int64_t offset = 0;
+        auto gathered = expert_flat.index_select(0, sort_order);
+        auto local_slot_output = at::zeros_like(flat);
+
+        // The grouped single-rank path knows that every sorted row is local
+        // and therefore needs no host-visible expert counts. Materialize the
+        // CPU copy lazily only for EP slicing or the legacy per-expert loop.
+        at::Tensor counts_cpu;
+        auto get_counts_cpu = [&]() -> const at::Tensor& {
+            if (!counts_cpu.defined()) {
+                counts_cpu = counts.to(at::TensorOptions().device(at::kCPU));
+            }
+            return counts_cpu;
+        };
+
+#if RUSTRAIN_HAS_ATEN_GROUPED_MM
+        // PyTorch 2.12+ exposes the same CUTLASS grouped-GEMM primitive used
+        // by its native MoE path. It consumes sorted token rows plus cumulative
+        // expert offsets, eliminating one GEMM launch per local expert. Older
+        // libtorch builds compile the fallback loop below.
+        bool grouped_lora_compatible = true;
+        if (expert_lora.gate_up_a) {
+            grouped_lora_compatible = expert_lora.gate_up_a->size(1) % 8 == 0;
+        }
+        if (expert_lora.down_a) {
+            grouped_lora_compatible = grouped_lora_compatible &&
+                expert_lora.down_a->size(1) % 8 == 0;
+        }
+        const bool use_grouped_mm = !env_enabled("QWEN36_DISABLE_GROUPED_MM") &&
+            compute_type == at::kBFloat16 && grouped_lora_compatible &&
+            hidden_dim % 8 == 0 && local_intermediate % 8 == 0;
+        if (use_grouped_mm) {
+            const bool owns_all_experts = expert_start == 0 &&
+                expert_count == counts.size(0);
+            const int64_t local_start = owns_all_experts || expert_start == 0
+                ? 0
+                : get_counts_cpu().narrow(0, 0, expert_start).sum().item<int64_t>();
+            const int64_t local_tokens = owns_all_experts
+                ? gathered.size(0)
+                : get_counts_cpu()
+                    .narrow(0, expert_start, expert_count).sum().item<int64_t>();
+            if (local_tokens > 0) {
+                if (env_enabled("QWEN36_REPORT_GROUPED_MM")) {
+                    std::fprintf(
+                        stderr,
+                        "[q36_moe] grouped_mm experts=%ld tokens=%ld hidden=%ld intermediate=%ld\n",
+                        static_cast<long>(expert_count),
+                        static_cast<long>(local_tokens),
+                        static_cast<long>(hidden_dim),
+                        static_cast<long>(local_intermediate));
+                }
+                auto selected = gathered.narrow(0, local_start, local_tokens);
+                auto token_indices = sort_order.narrow(
+                    0, local_start, local_tokens);
+                auto local_expert_indices = sorted_indices.narrow(
+                    0, local_start, local_tokens) - expert_start;
+                auto offsets = counts.narrow(0, expert_start, expert_count)
+                    .cumsum(0).to(at::kInt);
+                auto gu = at::_grouped_mm(
+                    selected, experts_gate_up.transpose(1, 2), offsets);
+                if (expert_lora.gate_up_a && expert_lora.gate_up_b) {
+                    auto lora_input = expert_tp
+                        ? selected : tp_copy_lora_input(training_ctx, selected);
+                    auto low_rank = at::_grouped_mm(
+                        lora_input, expert_lora.gate_up_a->transpose(1, 2), offsets);
+                    auto delta = at::_grouped_mm(
+                        low_rank, expert_lora.gate_up_b->transpose(1, 2), offsets);
+                    auto scaled = delta * expert_lora.scaling;
+                    gu = gu + (expert_tp
+                        ? scaled
+                        : tp_allreduce_lora_delta(training_ctx, scaled));
+                }
+                if (expert_gate_up_lora) {
+                    gu = gu + dynamic_expert_lora_delta(
+                        training_ctx, selected, token_indices, local_expert_indices,
+                        batch, seq, expert_gate_up_lora);
+                }
+                auto activated = fused_swiglu_op(
+                    gu.narrow(-1, 0, local_intermediate),
+                    gu.narrow(-1, local_intermediate, local_intermediate), 0.0);
+                auto expert_out = at::_grouped_mm(
+                    activated, experts_down.transpose(1, 2), offsets);
+                if (expert_lora.down_a && expert_lora.down_b) {
+                    auto lora_input = expert_tp
+                        ? activated : tp_copy_lora_input(training_ctx, activated);
+                    auto low_rank = at::_grouped_mm(
+                        lora_input, expert_lora.down_a->transpose(1, 2), offsets);
+                    auto delta = at::_grouped_mm(
+                        low_rank, expert_lora.down_b->transpose(1, 2), offsets);
+                    auto scaled = delta * expert_lora.scaling;
+                    expert_out = expert_out + (expert_tp
+                        ? scaled
+                        : tp_allreduce_lora_delta(training_ctx, scaled));
+                }
+                if (expert_down_lora) {
+                    expert_out = expert_out + dynamic_expert_lora_delta(
+                        training_ctx, activated, token_indices, local_expert_indices,
+                        batch, seq, expert_down_lora);
+                }
+                local_slot_output = local_slot_output.index_add(
+                    0, token_indices, expert_out);
+            }
+            local_slot_output = tp_allreduce_base_mlp(
+                training_ctx, local_slot_output);
+            routed_output = routed_output + local_slot_output *
+                expert_weights.unsqueeze(-1);
+            continue;
+        }
+#endif
+
+        // Process each expert's contiguous token slice. `gathered` contains
+        // all global experts in sorted order, so rank>0 must skip tokens for
+        // experts owned by lower ranks before taking its local range.
+        int64_t offset = expert_start > 0
+            ? get_counts_cpu().narrow(0, 0, expert_start).sum().item<int64_t>()
+            : 0;
         for (int64_t e_local = 0; e_local < expert_count; e_local++) {
             int64_t e_global = expert_start + e_local;
-            int64_t n_tokens = counts.index({e_global}).item<int64_t>();
-            if (n_tokens == 0) continue;
-
-            auto selected = gathered.narrow(0, offset, n_tokens);  // zero-copy view!
-            auto egu = experts_gate_up.select(0, e_local);
-            auto ed = experts_down.select(0, e_local);
-            auto gu = at::matmul(selected, egu.t());
-            auto gate_part = gu.narrow(-1, 0, intermediate);
-            auto up_part = gu.narrow(-1, intermediate, intermediate);
-            auto expert_out = at::matmul(fused_swiglu_op(gate_part, up_part, 0.0), ed.t());
-            auto weights = gathered_weights.narrow(0, offset, n_tokens);
-            auto token_indices = sort_order.narrow(0, offset, n_tokens);
-            routed_output = routed_output.index_add_(0, token_indices, expert_out * weights);
+            int64_t n_tokens = get_counts_cpu().index({e_global}).item<int64_t>();
+            if (n_tokens > 0) {
+                auto selected = gathered.narrow(0, offset, n_tokens);  // zero-copy view!
+                auto token_indices = sort_order.narrow(0, offset, n_tokens);
+                auto local_expert_indices = sorted_indices.narrow(0, offset, n_tokens)
+                    - expert_start;
+                auto egu = experts_gate_up.select(0, e_local);
+                auto ed = experts_down.select(0, e_local);
+                auto gu = at::matmul(selected, egu.t());
+                if (expert_lora.gate_up_a && expert_lora.gate_up_b) {
+                    auto a = expert_lora.gate_up_a->select(0, e_local);
+                    auto b = expert_lora.gate_up_b->select(0, e_local);
+                    auto lora_input = expert_tp
+                        ? selected : tp_copy_lora_input(training_ctx, selected);
+                    auto delta = at::matmul(at::matmul(lora_input, a.t()), b.t())
+                        * expert_lora.scaling;
+                    gu = gu + (expert_tp
+                        ? delta
+                        : tp_allreduce_lora_delta(training_ctx, delta));
+                }
+                if (expert_gate_up_lora) {
+                    gu = gu + dynamic_expert_lora_delta(
+                        training_ctx, selected, token_indices, local_expert_indices,
+                        batch, seq, expert_gate_up_lora);
+                }
+                auto gate_part = gu.narrow(-1, 0, local_intermediate);
+                auto up_part = gu.narrow(
+                    -1, local_intermediate, local_intermediate);
+                auto activated = fused_swiglu_op(gate_part, up_part, 0.0);
+                auto expert_out = at::matmul(activated, ed.t());
+                if (expert_lora.down_a && expert_lora.down_b) {
+                    auto a = expert_lora.down_a->select(0, e_local);
+                    auto b = expert_lora.down_b->select(0, e_local);
+                    auto lora_input = expert_tp
+                        ? activated : tp_copy_lora_input(training_ctx, activated);
+                    auto delta = at::matmul(at::matmul(lora_input, a.t()), b.t())
+                        * expert_lora.scaling;
+                    expert_out = expert_out + (expert_tp
+                        ? delta
+                        : tp_allreduce_lora_delta(training_ctx, delta));
+                }
+                if (expert_down_lora) {
+                    expert_out = expert_out + dynamic_expert_lora_delta(
+                        training_ctx, activated, token_indices, local_expert_indices,
+                        batch, seq, expert_down_lora);
+                }
+                local_slot_output = local_slot_output.index_add(
+                    0, token_indices, expert_out);
+            }
             offset += n_tokens;
         }
+        local_slot_output = tp_allreduce_base_mlp(
+            training_ctx, local_slot_output);
+        routed_output = routed_output + local_slot_output *
+            expert_weights.unsqueeze(-1);
+    }
+
+    // A rank can receive no tokens for any of its local experts. Keep a
+    // zero-valued dependency on routed-expert LoRA tensors so autograd still
+    // produces defined zero gradients and every rank reaches the same NCCL
+    // collectives. This changes only the graph, not the routed output values.
+    if (!use_a2a && !routed_output.requires_grad()) {
+        at::Tensor graph_anchor;
+        auto include_anchor = [&](const at::Tensor* tensor) {
+            if (!tensor || !tensor->defined() || !tensor->requires_grad()) return;
+            auto contribution = tensor->sum().to(routed_output.scalar_type());
+            graph_anchor = graph_anchor.defined()
+                ? graph_anchor + contribution
+                : contribution;
+        };
+        include_anchor(expert_lora.gate_up_a);
+        include_anchor(expert_lora.gate_up_b);
+        include_anchor(expert_lora.down_a);
+        include_anchor(expert_lora.down_b);
+        if (expert_gate_up_lora) {
+            include_anchor(&expert_gate_up_lora->a_stack);
+            include_anchor(&expert_gate_up_lora->b_stack);
+        }
+        if (expert_down_lora) {
+            include_anchor(&expert_down_lora->a_stack);
+            include_anchor(&expert_down_lora->b_stack);
+        }
+        if (graph_anchor.defined()) {
+            routed_output = routed_output + graph_anchor * 0.0;
+        }
     }
 
     // EP all-reduce via NcclAllReduceFunction — custom autograd Function.
-    if (nccl_comm_v) {
+    if (nccl_comm_v && !use_a2a) {
         auto nccl_comm = reinterpret_cast<ncclComm_t>(nccl_comm_v);
-        routed_output = NcclAllReduceFunction::apply(routed_output, (int64_t)nccl_comm);
+        routed_output = NcclAllReduceFunction::apply(
+            routed_output, (int64_t)nccl_comm,
+            (int64_t)reinterpret_cast<uintptr_t>(nccl_stream_v));
     }
 
-    // Shared expert (same as before, with fused SwiGLU)
-    auto shared_gate = at::matmul(flat, shared_gate_proj.t());
-    auto shared_up = at::matmul(flat, shared_up_proj.t());
-    auto shared_out = at::matmul(fused_swiglu_op(shared_gate, shared_up, 0.0), shared_down_proj.t());
-    auto seg = at::sigmoid(at::matmul(flat, shared_expert_gate_w.t())).to(compute_type);
-    shared_out = (shared_out * seg).to(compute_type);
+    if (!shared_overlap_active) {
+        shared_out = compute_shared_expert();
+    } else {
+        const auto current_stream =
+            c10::cuda::getCurrentCUDAStream(device.index()).stream();
+        TORCH_CHECK(cudaStreamWaitEvent(
+                current_stream, shared_events.done, 0) == cudaSuccess,
+            "failed to fence compute stream before shared-expert combine");
+    }
 
     // Debug: dump routed_output AFTER loop
     if (getenv("QWEN36_DUMP_MOE")) {
@@ -772,28 +3490,121 @@ static inline int64_t weight_count_for_layer(const LayerConfig& cfg) {
     return 2 + attn_w + mlp_w;
 }
 
-static at::Tensor forward_single_layer(
-    TrainingContext* ctx, const at::Tensor& hidden, at::Tensor** w, const LayerConfig* cfg,
-    int64_t layer_idx, at::ScalarType kind,
-    const at::Tensor& attention_mask, bool use_batched = false
-) {
-    auto input_norm = *w[0];
-    auto post_norm = *w[1];
-    auto attn_input = rms_norm(hidden, input_norm, cfg->rms_eps);
-    bool is_moe = (cfg->num_experts > 0);
+enum class LoraSegment : uint8_t { Attention, Mlp };
 
-    at::Tensor attn_output;
-    if (cfg->layer_type == 0) {
-        // Full attention
-        auto q_proj = *w[2], q_norm = *w[3], k_proj = *w[4], k_norm = *w[5], v_proj = *w[6], o_proj = *w[7];
-        if (use_batched) {
+static bool is_mlp_lora_target(const std::string& name) {
+    return name == "gate_proj" || name == "up_proj" || name == "down_proj" ||
+        name == "shared_gate_proj" || name == "shared_up_proj" ||
+        name == "shared_down_proj" || name == "experts_gate_up_proj" ||
+        name == "experts_down_proj";
+}
+
+struct LoraProjectionSpec {
+    const char* name;
+    int64_t weight_index;
+    LoraSegment segment;
+    bool grouped_expert;
+};
+
+struct LoraProjectionTable {
+    std::array<LoraProjectionSpec, 10> entries;
+    int64_t count = 0;
+
+    void add(const char* name, int64_t weight_index, LoraSegment segment) {
+        TORCH_CHECK(count < (int64_t)entries.size(), "too many LoRA projections in layer");
+        entries[count++] = {name, weight_index, segment, false};
+    }
+
+    void add_grouped_expert(const char* name, int64_t weight_index) {
+        TORCH_CHECK(count < (int64_t)entries.size(), "too many LoRA projections in layer");
+        entries[count++] = {name, weight_index, LoraSegment::Mlp, true};
+    }
+};
+
+static LoraProjectionTable lora_projection_table(const LayerConfig& cfg) {
+    LoraProjectionTable table;
+    if (cfg.layer_type == 0) {
+        table.add("q_proj", 2, LoraSegment::Attention);
+        table.add("k_proj", 4, LoraSegment::Attention);
+        table.add("v_proj", 6, LoraSegment::Attention);
+        table.add("o_proj", 7, LoraSegment::Attention);
+    } else {
+        table.add("in_proj_qkv", 2, LoraSegment::Attention);
+        table.add("in_proj_z", 3, LoraSegment::Attention);
+        table.add("in_proj_a", 4, LoraSegment::Attention);
+        table.add("in_proj_b", 5, LoraSegment::Attention);
+        table.add("out_proj", 10, LoraSegment::Attention);
+    }
+
+    const int64_t mlp_start = cfg.layer_type == 0 ? 8 : 11;
+    if (cfg.num_experts > 0) {
+        table.add("shared_gate_proj", mlp_start + 2, LoraSegment::Mlp);
+        table.add("shared_up_proj", mlp_start + 3, LoraSegment::Mlp);
+        table.add("shared_down_proj", mlp_start + 4, LoraSegment::Mlp);
+        table.add_grouped_expert("experts_gate_up_proj", mlp_start + 5);
+        table.add_grouped_expert("experts_down_proj", mlp_start + 6);
+    } else {
+        table.add("gate_proj", mlp_start, LoraSegment::Mlp);
+        table.add("up_proj", mlp_start + 1, LoraSegment::Mlp);
+        table.add("down_proj", mlp_start + 2, LoraSegment::Mlp);
+    }
+    return table;
+}
+
+static inline int64_t lora_pair_count(const LayerConfig& cfg) {
+    return lora_projection_table(cfg).count;
+}
+
+static int64_t lora_pair_index(const LayerConfig& cfg, const char* name) {
+    auto table = lora_projection_table(cfg);
+    for (int64_t i = 0; i < table.count; ++i) {
+        if (std::strcmp(table.entries[i].name, name) == 0) return i;
+    }
+    return -1;
+}
+
+static RoutedExpertLora routed_expert_lora(
+    TrainingContext* ctx, int64_t layer_idx, const LayerConfig& cfg);
+static at::Tensor tp_allreduce_base_mlp(
+    TrainingContext* ctx, const at::Tensor& local_output);
+static at::Tensor tp_copy_base_mlp_input(
+    TrainingContext* ctx, const at::Tensor& input);
+static at::Tensor tp_allreduce_base_attention(
+    TrainingContext* ctx, const at::Tensor& local_output);
+static at::Tensor tp_copy_base_attention_input(
+    TrainingContext* ctx, const at::Tensor& input);
+static bool base_tp_attention_enabled(const TrainingContext* ctx);
+static at::Tensor fused_full_attention_qkv_weight(
+    TrainingContext* ctx, int64_t layer_idx);
+
+static at::Tensor forward_single_layer(
+    TrainingContext* ctx, const at::Tensor& hidden, at::Tensor** w, const LayerConfig* cfg,
+    int64_t layer_idx, at::ScalarType kind,
+    const at::Tensor& attention_mask, const at::Tensor& attention_lengths,
+    bool use_batched = false
+) {
+    auto input_norm = *w[0];
+    auto post_norm = *w[1];
+    auto attn_input = rms_norm(hidden, input_norm, cfg->rms_eps);
+    bool is_moe = (cfg->num_experts > 0);
+
+    at::Tensor attn_output;
+    if (cfg->layer_type == 0) {
+        // Full attention
+        auto q_proj = *w[2], q_norm = *w[3], k_proj = *w[4], k_norm = *w[5], v_proj = *w[6], o_proj = *w[7];
+        const at::Tensor fused_qkv = use_batched
+            ? fused_full_attention_qkv_weight(ctx, layer_idx)
+            : at::Tensor();
+        TORCH_CHECK(!base_tp_attention_enabled(ctx) || use_batched,
+            "base full-attention TP requires the activation-level LoRA path");
+        if (use_batched) {
             // Activation-level LoRA: pass base weights, apply delta inside attention
             attn_output = full_attention_batched(
                 ctx, attn_input, layer_idx,
                 q_proj, q_norm, k_proj, k_norm, v_proj, o_proj,
                 cfg->num_heads, cfg->num_kv_heads, cfg->head_dim,
                 cfg->partial_rotary_factor, cfg->rope_theta, cfg->rms_eps, kind,
-                attention_mask);
+                attention_mask, fused_qkv);
         } else {
             // Weight-level LoRA (legacy: modify weights, single forward)
             q_proj = apply_multi_lora(ctx, layer_idx, 0, q_proj);
@@ -807,44 +3618,114 @@ static at::Tensor forward_single_layer(
         }
         auto post_attn = rms_norm(hidden + attn_output, post_norm, cfg->rms_eps);
         if (is_moe) {
-            auto mlp_out = moe_forward(cfg->nccl_comm, cfg->nccl_stream, post_attn,
-                *w[8], *w[9], *w[10], *w[11], *w[12], *w[13], *w[14],
+            const int64_t shared_gate_pair = lora_pair_index(*cfg, "shared_gate_proj");
+            const int64_t shared_up_pair = lora_pair_index(*cfg, "shared_up_proj");
+            const int64_t shared_down_pair = lora_pair_index(*cfg, "shared_down_proj");
+            const int64_t expert_gate_up_pair = lora_pair_index(*cfg, "experts_gate_up_proj");
+            const int64_t expert_down_pair = lora_pair_index(*cfg, "experts_down_proj");
+            auto shared_gate = use_batched ? *w[10]
+                : apply_multi_lora(ctx, layer_idx, shared_gate_pair, *w[10]);
+            auto shared_up = use_batched ? *w[11]
+                : apply_multi_lora(ctx, layer_idx, shared_up_pair, *w[11]);
+            auto shared_down = use_batched ? *w[12]
+                : apply_multi_lora(ctx, layer_idx, shared_down_pair, *w[12]);
+            auto expert_lora = routed_expert_lora(ctx, layer_idx, *cfg);
+            auto mlp_out = moe_forward(ctx, layer_idx,
+                cfg->nccl_comm, cfg->nccl_stream, post_attn,
+                *w[8], *w[9], shared_gate, shared_up, shared_down, *w[13], *w[14],
+                expert_lora,
                 cfg->num_experts, cfg->top_k, cfg->moe_intermediate,
-                cfg->norm_topk_prob != 0, cfg->expert_start, cfg->expert_count, kind);
+                cfg->norm_topk_prob != 0, cfg->expert_start, cfg->expert_count,
+                kind, use_batched,
+                use_batched ? lora_batch_entry(ctx, layer_idx, shared_gate_pair) : nullptr,
+                use_batched ? lora_batch_entry(ctx, layer_idx, shared_up_pair) : nullptr,
+                use_batched ? lora_batch_entry(ctx, layer_idx, shared_down_pair) : nullptr,
+                use_batched ? lora_batch_entry(ctx, layer_idx, expert_gate_up_pair) : nullptr,
+                use_batched ? lora_batch_entry(ctx, layer_idx, expert_down_pair) : nullptr);
             return hidden + attn_output + mlp_out;
         } else {
-            auto mlp_out = dense_mlp_forward(post_attn, *w[8], *w[9], *w[10], kind);
+            if (use_batched) {
+                auto mlp_out = dense_mlp_forward_batched(
+                    ctx, layer_idx, post_attn, *w[8], *w[9], *w[10], kind, cfg);
+                return hidden + attn_output + mlp_out;
+            }
+            auto gate = apply_multi_lora(ctx, layer_idx,
+                lora_pair_index(*cfg, "gate_proj"), *w[8]);
+            auto up = apply_multi_lora(ctx, layer_idx,
+                lora_pair_index(*cfg, "up_proj"), *w[9]);
+            auto down = apply_multi_lora(ctx, layer_idx,
+                lora_pair_index(*cfg, "down_proj"), *w[10]);
+            auto mlp_input = tp_copy_base_mlp_input(ctx, post_attn);
+            auto mlp_out = tp_allreduce_base_mlp(
+                ctx, dense_mlp_forward(mlp_input, gate, up, down, kind));
             return hidden + attn_output + mlp_out;
         }
     } else {
         // Linear attention
         auto in_proj_qkv = *w[2], in_proj_z = *w[3], in_proj_a = *w[4], in_proj_b = *w[5];
         auto a_log = *w[6], dt_bias = *w[7], conv1d_w = *w[8], norm_w = *w[9], out_proj = *w[10];
+        TORCH_CHECK(!base_tp_attention_enabled(ctx) || use_batched,
+            "base linear-attention TP requires the activation-level LoRA path");
         if (use_batched) {
             attn_output = linear_attention_batched(
                 ctx, attn_input, layer_idx,
                 in_proj_qkv, in_proj_z, in_proj_a, in_proj_b,
                 a_log, dt_bias, conv1d_w, norm_w, out_proj,
                 cfg->num_k_heads, cfg->key_dim, cfg->num_v_heads, cfg->val_dim,
-                cfg->conv_kernel, cfg->rms_eps, kind);
+                cfg->conv_kernel, cfg->rms_eps, kind, attention_lengths);
         } else {
             in_proj_qkv = apply_multi_lora(ctx, layer_idx, 0, in_proj_qkv);
             in_proj_z = apply_multi_lora(ctx, layer_idx, 1, in_proj_z);
-            out_proj = apply_multi_lora(ctx, layer_idx, 2, out_proj);
+            in_proj_a = apply_multi_lora(ctx, layer_idx, 2, in_proj_a);
+            in_proj_b = apply_multi_lora(ctx, layer_idx, 3, in_proj_b);
+            out_proj = apply_multi_lora(ctx, layer_idx, 4, out_proj);
             attn_output = linear_attention(attn_input, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b,
                 a_log, dt_bias, conv1d_w, norm_w, out_proj,
                 cfg->num_k_heads, cfg->key_dim, cfg->num_v_heads, cfg->val_dim,
-                cfg->conv_kernel, cfg->rms_eps, kind);
+                cfg->conv_kernel, cfg->rms_eps, kind, attention_lengths);
         }
         auto post_attn = rms_norm(hidden + attn_output, post_norm, cfg->rms_eps);
         if (is_moe) {
-            auto mlp_out = moe_forward(cfg->nccl_comm, cfg->nccl_stream, post_attn,
-                *w[11], *w[12], *w[13], *w[14], *w[15], *w[16], *w[17],
+            const int64_t shared_gate_pair = lora_pair_index(*cfg, "shared_gate_proj");
+            const int64_t shared_up_pair = lora_pair_index(*cfg, "shared_up_proj");
+            const int64_t shared_down_pair = lora_pair_index(*cfg, "shared_down_proj");
+            const int64_t expert_gate_up_pair = lora_pair_index(*cfg, "experts_gate_up_proj");
+            const int64_t expert_down_pair = lora_pair_index(*cfg, "experts_down_proj");
+            auto shared_gate = use_batched ? *w[13]
+                : apply_multi_lora(ctx, layer_idx, shared_gate_pair, *w[13]);
+            auto shared_up = use_batched ? *w[14]
+                : apply_multi_lora(ctx, layer_idx, shared_up_pair, *w[14]);
+            auto shared_down = use_batched ? *w[15]
+                : apply_multi_lora(ctx, layer_idx, shared_down_pair, *w[15]);
+            auto expert_lora = routed_expert_lora(ctx, layer_idx, *cfg);
+            auto mlp_out = moe_forward(ctx, layer_idx,
+                cfg->nccl_comm, cfg->nccl_stream, post_attn,
+                *w[11], *w[12], shared_gate, shared_up, shared_down, *w[16], *w[17],
+                expert_lora,
                 cfg->num_experts, cfg->top_k, cfg->moe_intermediate,
-                cfg->norm_topk_prob != 0, cfg->expert_start, cfg->expert_count, kind);
+                cfg->norm_topk_prob != 0, cfg->expert_start, cfg->expert_count,
+                kind, use_batched,
+                use_batched ? lora_batch_entry(ctx, layer_idx, shared_gate_pair) : nullptr,
+                use_batched ? lora_batch_entry(ctx, layer_idx, shared_up_pair) : nullptr,
+                use_batched ? lora_batch_entry(ctx, layer_idx, shared_down_pair) : nullptr,
+                use_batched ? lora_batch_entry(ctx, layer_idx, expert_gate_up_pair) : nullptr,
+                use_batched ? lora_batch_entry(ctx, layer_idx, expert_down_pair) : nullptr);
             return hidden + attn_output + mlp_out;
         } else {
-            auto mlp_out = dense_mlp_forward(post_attn, *w[11], *w[12], *w[13], kind);
+            if (use_batched) {
+                auto mlp_out = dense_mlp_forward_batched(
+                    ctx, layer_idx, post_attn, *w[11], *w[12], *w[13], kind, cfg);
+                return hidden + attn_output + mlp_out;
+            }
+            auto gate = apply_multi_lora(ctx, layer_idx,
+                lora_pair_index(*cfg, "gate_proj"), *w[11]);
+            auto up = apply_multi_lora(ctx, layer_idx,
+                lora_pair_index(*cfg, "up_proj"), *w[12]);
+            auto down = apply_multi_lora(ctx, layer_idx,
+                lora_pair_index(*cfg, "down_proj"), *w[13]);
+            auto mlp_input = tp_copy_base_mlp_input(ctx, post_attn);
+            auto mlp_out = tp_allreduce_base_mlp(
+                ctx, dense_mlp_forward(mlp_input, gate, up, down, kind));
             return hidden + attn_output + mlp_out;
         }
     }
@@ -859,8 +3740,18 @@ struct AdamDevBuffers {
     at::Tensor grads_buf;    // [max_n] kLong
     at::Tensor m_buf;        // [max_n] kLong — float* stored as int64_t
     at::Tensor v_buf;        // [max_n] kLong
+    at::Tensor dst_params_buf;  // [max_n] kLong
+    at::Tensor dst_m_buf;       // [max_n] kLong
+    at::Tensor dst_v_buf;       // [max_n] kLong
     at::Tensor sizes_buf;    // [max_n] kInt
+    at::Tensor lr_buf;       // [max_n] kFloat
+    at::Tensor eps_buf;      // [max_n] kFloat
+    at::Tensor beta1_buf;    // [max_n] kFloat
+    at::Tensor beta2_buf;    // [max_n] kFloat
+    at::Tensor groups_buf;   // [max_n] kInt
+    at::Tensor clip_norm_squares;  // [max_groups] kFloat
     int capacity = 0;
+    int group_capacity = 0;
 
     void ensure(int n, const at::Tensor& ref) {
         if (n <= capacity) return;
@@ -869,36 +3760,181 @@ struct AdamDevBuffers {
         grads_buf  = at::empty({n}, at::TensorOptions().dtype(at::kLong).device(dev));
         m_buf      = at::empty({n}, at::TensorOptions().dtype(at::kLong).device(dev));
         v_buf      = at::empty({n}, at::TensorOptions().dtype(at::kLong).device(dev));
+        dst_params_buf = at::empty({n}, at::TensorOptions().dtype(at::kLong).device(dev));
+        dst_m_buf = at::empty({n}, at::TensorOptions().dtype(at::kLong).device(dev));
+        dst_v_buf = at::empty({n}, at::TensorOptions().dtype(at::kLong).device(dev));
         sizes_buf  = at::empty({n}, at::TensorOptions().dtype(at::kInt).device(dev));
+        lr_buf     = at::empty({n}, at::TensorOptions().dtype(at::kFloat).device(dev));
+        eps_buf    = at::empty({n}, at::TensorOptions().dtype(at::kFloat).device(dev));
+        beta1_buf  = at::empty({n}, at::TensorOptions().dtype(at::kFloat).device(dev));
+        beta2_buf  = at::empty({n}, at::TensorOptions().dtype(at::kFloat).device(dev));
+        groups_buf = at::empty({n}, at::TensorOptions().dtype(at::kInt).device(dev));
         capacity = n;
     }
+
+    void ensure_groups(int n, const at::Tensor& ref) {
+        if (n <= group_capacity) return;
+        clip_norm_squares = at::empty(
+            {n}, at::TensorOptions().dtype(at::kFloat).device(ref.device()));
+        group_capacity = n;
+    }
+};
+
+struct PipelineWindowSlot {
+    at::Tensor input_ids;
+    at::Tensor target_mask;
+    at::Tensor attention_mask;
+    at::Tensor stage_input;
+    at::Tensor stage_output;
+    double gradient_scale = 0.0;
+    double token_weight = 0.0;
+    at::Tensor row_token_weights;
+    bool forward_done = false;
+    bool backward_done = false;
+};
+
+struct PipelineWindowState {
+    bool active = false;
+    int64_t window_id = -1;
+    int64_t num_microbatches = 0;
+    int32_t schedule = 0;
+    int32_t num_chunks = 1;
+    int32_t flags = 0;
+    bool dynamic_lora = false;
+    bool previous_pad_heterogeneous_lora_batch = false;
+    bool padding_mode_overridden = false;
+    std::vector<size_t> selected_adapter_indices;
+    int64_t next_forward = 0;
+    int64_t next_backward = 0;
+    double total_loss = 0.0;
+    double total_loss_weight = 0.0;
+    at::Tensor adapter_token_counts;
+    at::Tensor loss_numerators;
+    at::Tensor loss_token_counts;
+    at::Tensor normalization_mask;
+    at::Tensor pending_backward_send;
+    int64_t pending_backward_mb = -1;
+    int64_t batch_size = -1;
+    int64_t sequence_length = -1;
+    int32_t input_dtype = -1;
+    int32_t target_dtype = -1;
+    int32_t attention_present = -1;
+    int32_t attention_dtype = -1;
+    bool local_error = false;
+    std::map<int64_t, PipelineWindowSlot> slots;
 };
 
 struct TrainingContext {
+    // All ranks create contexts in the same order. This sequence isolates
+    // per-context rendezvous files when a worker reuses one NCCL process.
+    int64_t context_sequence = 0;
     // Model weights (frozen, no grad) — pointers to external tensors
     std::vector<at::Tensor*> weight_ptrs;  // flat array, 15 or 18 per layer
     std::vector<at::Tensor*> embed_ptr;
     std::vector<at::Tensor*> final_norm_ptr;
     std::vector<at::Tensor*> lm_head_ptr;
     std::vector<LayerConfig> layer_configs;
+    // Frozen GDN alpha/beta projections are tiny and launch-bound. Keep a
+    // context-owned concatenation so the activation-level LoRA path computes
+    // both with one GEMM without duplicating the large QKV/Z weights.
+    std::vector<at::Tensor> fused_gdn_ab_weights;
+    // Optional full-attention base QKV concatenation. This is opt-in because
+    // it duplicates frozen Q/K/V storage, but it reduces three base GEMMs to
+    // one for activation-level LoRA and TP-local attention.
+    std::vector<at::Tensor> fused_full_attention_qkv_weights;
+    // Optional frozen dense/shared SwiGLU FC1 concatenation. Dynamic LoRA
+    // deltas remain separate per tenant and are added after the single base
+    // projection is split back into gate/up activations.
+    std::vector<at::Tensor> fused_mlp_fc1_weights;
     int64_t num_layers;
 
     // Attention mask [batch, seq] — 1 for real tokens, 0 for padding
     at::Tensor attention_mask;
+    // The Rust training loop passes input IDs directly to the native entry
+    // point. When no explicit mask is supplied, native code can derive the
+    // padding mask without a Rust-side GPU comparison/kernel launch.
+    int64_t pad_token_id = -1;
+    // Strict-right-padding valid lengths for GDN. Kept on device so every
+    // layer can pass it into the persistent recurrent kernel without a host
+    // sync or repeated mask reduction.
+    at::Tensor attention_lengths;
+
+    // Optional host staging for the shared-memory i64 ingress. The IPC slab
+    // is pageable and may be reused as soon as the native call returns, so
+    // each logical input owns a separate pinned buffer before an async H2D
+    // copy is enqueued. The buffers are resized only when the request shape
+    // changes and are kept context-owned until the next synchronous request.
+    at::Tensor host_i64_input_staging;
+    at::Tensor host_i64_target_staging;
+    at::Tensor host_i64_attention_staging;
 
     // ── Multi-LoRA adapter registry ──
     struct LoRAAdapter {
         int64_t id;
         int64_t rank;
+        // Each tenant owns an independent Adam bias-correction clock.
+        int64_t optimizer_step = 0;
+        // Dynamic tenants own complete Adam hyperparameters while sharing one
+        // fused launch through per-tensor scalar buffers.
+        double optimizer_lr = 0.0;
+        double optimizer_beta1 = 0.9;
+        double optimizer_beta2 = 0.999;
+        double optimizer_eps = 1e-8;
         double alpha;
+        bool all_target_layers = true;
+        // External identity remains global across pipeline stages. The local
+        // set below is used only for stage-owned parameter maps and compute.
+        std::set<int64_t> global_target_layers;
         std::set<int64_t> target_layers;
         std::set<std::string> target_modules;
         std::map<int64_t, std::vector<std::pair<at::Tensor, at::Tensor>>> params;
         std::map<int64_t, std::vector<std::array<at::Tensor, 4>>> adam_state;
+        // Checkpoint snapshots are host-resident and versioned by the
+        // tenant-local Adam clock. They never participate in the training
+        // transaction; a successful step invalidates them by advancing
+        // optimizer_step, while rollback makes the previous snapshot current
+        // again without copying.
+        std::map<int64_t, std::vector<std::array<at::Tensor, 4>>>
+            adam_host_state;
+        int64_t adam_host_step = -1;
+        uint64_t adam_last_used_tick = 0;
+        // Reusable out-of-place Adam destinations. Each entry stores
+        // {param_a, m_a, v_a, param_b, m_b, v_b}; a successful transaction
+        // swaps these handles with the live registry without allocating.
+        std::map<int64_t, std::vector<std::array<at::Tensor, 6>>> adam_shadow;
+        // When pooled shadows are enabled, each active pair borrows one
+        // context-owned slot for the duration of a logical transaction.
+        // Inactive adapters retain only these small integer lease records.
+        std::map<int64_t, std::vector<int64_t>> adam_shadow_slots;
+        // Gradients are harvested after each backward into FP32 tensors. The
+        // accumulator tensors are intentionally shared by value when chunk
+        // registry guards copy adapters, so their contents survive restore.
+        std::map<int64_t, std::vector<std::array<at::Tensor, 2>>> grad_accum;
+        at::Tensor grad_slab;
     };
 
     std::vector<LoRAAdapter> adapters;
+    struct DynamicAdamShadowPoolEntry {
+        std::array<at::Tensor, 6> tensors;
+        bool in_use = false;
+    };
+    // Transaction destinations scale with the maximum simultaneously active
+    // adapter layouts instead of every tenant that has ever trained.
+    std::vector<DynamicAdamShadowPoolEntry> dynamic_adam_shadow_pool;
+    // Sorted external ID -> canonical registry slot. Temporary selected/chunk
+    // registries explicitly invalidate this index; their adapter order is not
+    // the canonical order represented here.
+    std::vector<std::pair<int64_t, size_t>> adapter_id_index;
+    bool adapter_id_index_valid = true;
     int64_t next_adapter_id = 0;
+    int64_t multi_lora_invocation = 0;
+    int64_t dynamic_finalizer_count = 0;
+    int64_t dynamic_adam_launch_count = 0;
+    uint64_t dynamic_adam_lru_tick = 0;
+    bool dynamic_adam_transaction_active = false;
+    bool restore_without_parameter_sync = false;
+    bool allow_heterogeneous_registration = false;
+    bool pad_heterogeneous_lora_batch = false;
 
     // LoRA cache: pre-concatenated A/B per (layer, module) pair
     // Invalidated when adapters change or after Adam update
@@ -908,17 +3944,24 @@ struct TrainingContext {
     // ── Batched Multi-LoRA (activation-level) ──
     // When active, replaces the weight-level lora_cache.
     // Stores per-(layer, module) stacked A/B tensors for batched B@(A@x) computation.
-    struct LoraBatchEntry {
-        at::Tensor a_stack;    // [N, rank, in]
-        at::Tensor b_stack;    // [N, out, rank]
-        at::Tensor scaling;    // [N, 1, 1] — per-adapter alpha/rank
-    };
     bool lora_batch_valid = false;
     int64_t lora_batch_n = 0;  // number of adapters in current batch
     std::map<int64_t, LoraBatchEntry> lora_batch_cache;
+    // Shared BF16 alpha/rank scale for the current activation batch. Building
+    // it once avoids one tiny CPU->GPU transfer per layer/module projection.
+    at::Tensor lora_batch_scaling;
+    std::vector<double> lora_batch_scaling_values;
+    int64_t lora_batch_projection_build_count = 0;
+    int64_t lora_batch_scaling_upload_count = 0;
     // Legacy single-LoRA (backward compat)
     std::vector<at::Tensor> lora_a;
     std::vector<at::Tensor> lora_b;
+    // Fixed-slot FP32 gradient accumulators. Inactive slots are undefined
+    // tensors to preserve the positional LoRA ABI without extra allocation.
+    std::vector<at::Tensor> grad_accum_a;
+    std::vector<at::Tensor> grad_accum_b;
+    at::Tensor fixed_grad_slab;
+    std::vector<uint8_t> lora_active;
     std::vector<int64_t> lora_layer_offset;
     double lora_scaling;
     std::vector<std::string> lora_names;
@@ -927,7 +3970,32 @@ struct TrainingContext {
     std::vector<at::Tensor> adam_m;
     std::vector<at::Tensor> adam_v;
     double lr, beta1, beta2, eps;
-    int64_t step_count;
+    // Zero disables clipping. Positive values apply one logical global norm
+    // to fixed LoRA and one independent global norm per dynamic tenant.
+    double max_grad_norm = 0.0;
+    // Switch-style router load-balancing loss. The forward attachment keeps
+    // routing probabilities bit-identical and injects the auxiliary gradient
+    // when the model graph is traversed.
+    double router_aux_loss_coef = 0.0;
+    at::Tensor router_aux_backward_scale;
+    bool collect_router_aux_loss = false;
+    at::Tensor router_aux_loss_value;
+    // Fixed-LoRA target identity is global even when the context owns only a
+    // pipeline stage. PP consensus hashes this metadata instead of local
+    // tensor slots, which legitimately differ between stages.
+    bool fixed_all_target_layers = true;
+    std::set<int64_t> fixed_target_layers_global;
+    std::set<std::string> fixed_target_modules;
+    // The fixed adapter's Adam bias-correction clock. Dynamic tenant updates
+    // must not advance it because every tenant has its own optimizer_step.
+    int64_t fixed_optimizer_step;
+    // A failed native call aborts the in-flight gradient window. This flag is
+    // consumed by the scoped accumulation guard on stack unwinding.
+    bool accumulation_active = false;
+    // Sum of (micro weight * supervised tokens) for the fixed-LoRA window.
+    // Fixed gradients are accumulated as token-weighted numerators and
+    // normalized exactly once at the optimizer boundary.
+    double accumulated_token_weight = 0.0;
 
     // Device buffer cache for multi-tensor fused Adam
     AdamDevBuffers adam_dev_bufs;
@@ -936,9 +4004,14 @@ struct TrainingContext {
     at::ScalarType compute_type;
     int64_t vocab_size;
     double rms_eps;
+    int64_t global_layer_start = 0;
+    int64_t global_num_layers = 0;
+    bool is_first_pipeline_stage = true;
+    bool is_last_pipeline_stage = true;
 
     // MTP weights (optional)
     bool has_mtp;
+    double mtp_loss_scale = 0.1;  // NVIDIA Megatron default
     at::Tensor *mtp_fc, *mtp_pre_fc_norm_emb, *mtp_pre_fc_norm_hidden, *mtp_norm;
     std::vector<at::Tensor*> mtp_layer_weights;
     std::vector<LayerConfig> mtp_layer_configs;
@@ -952,2708 +4025,13639 @@ struct TrainingContext {
     // Variable-size group ranges for selective checkpointing
     std::vector<std::pair<int64_t, int64_t>> group_ranges;
 
-    // NCCL for Expert Parallel all-reduce (nullptr if single-GPU)
+    // Expert-parallel communicator. Layer dispatch/combine and dense replica
+    // reductions use this axis; routed-expert parameters never reduce on it.
     ncclComm_t nccl_comm = nullptr;
-    cudaStream_t nccl_stream = nullptr;
+    // Use PyTorch's compute stream for NCCL — NOT a separate stream.
+    // Separate stream causes "invalid argument" because NCCL communicator
+    // is bound to a CUDA context, and PyTorch's caching allocator stream
+    // may be on a different context. Using the same stream ensures same context.
+    // We store nullptr for stream — moe_forward will use getCurrentCUDAStream(dev).
+    cudaStream_t nccl_stream = nullptr;  // nullptr = use default stream
+    // Stable context-owned compute stream for optional routed/shared expert
+    // overlap. A stream-pool lookup per layer retains one allocator workspace
+    // per pooled stream, so the runtime creates exactly one stream lazily.
+    cudaStream_t moe_shared_stream = nullptr;
+    // The variable-count exchange uses this stream only for the count
+    // all-gather and pinned D2H copy. Payload A2A retains its existing stream
+    // and fences, while send packing proceeds concurrently on the compute
+    // stream.
+    cudaStream_t moe_ep_metadata_stream = nullptr;
     int ep_world_size = 1;
     int ep_rank = 0;
+    bool expert_parallel = false;
+    // Expert-data-parallel communicator. Routed experts are replicated only
+    // across this axis, while dense parameters reduce across both EP and DP.
+    ncclComm_t dp_comm = nullptr;
+    cudaStream_t dp_stream = nullptr;
+    int dp_world_size = 1;
+    int dp_rank = 0;
+    bool data_parallel = false;
+    // LoRA-only tensor parallelism keeps frozen base weights replicated and
+    // shards the latent rank. Only the local LoRA delta uses this communicator.
+    ncclComm_t tp_comm = nullptr;
+    cudaStream_t tp_stream = nullptr;
+    int tp_world_size = 1;
+    int tp_rank = 0;
+    // TP2 sequence parallelism partitions hidden activations over sequence
+    // while frozen projection weights retain their existing TP shards.
+    bool sequence_parallel = false;
+    int64_t sequence_parallel_embedding_scatter_count = 0;
+    int64_t sequence_parallel_all_gather_count = 0;
+    int64_t sequence_parallel_reduce_scatter_count = 0;
+    int64_t sequence_parallel_loss_gather_count = 0;
+    int64_t sequence_parallel_last_local_sequence = 0;
+    // Context and pipeline communicators are initialized as part of the
+    // five-dimensional process grid. Model execution remains fail-closed
+    // until sequence partitioning and stage ownership are implemented.
+    ncclComm_t cp_comm = nullptr;
+    cudaStream_t cp_stream = nullptr;
+    int cp_world_size = 1;
+    int cp_rank = 0;
+    ncclComm_t pp_comm = nullptr;
+    cudaStream_t pp_stream = nullptr;
+    // PP control collectives use a duplicate communicator so shape/registry
+    // checks cannot be ordered against activation/gradient P2P traffic.
+    ncclComm_t pp_control_comm = nullptr;
+    int pp_world_size = 1;
+    int pp_rank = 0;
+    PipelineWindowState pp_window;
+    // Frozen dense SwiGLU TP: gate/up are output-sharded and down is
+    // input-sharded by the Rust weight loader.  The local row contribution
+    // is reduced over the TP communicator before the residual add.
+    bool base_tp_mlp = false;
+    // Frozen attention TP: full attention and GDN own disjoint head bundles;
+    // their output projections own the matching input columns.
+    bool base_tp_attention = false;
+    // Vocabulary parallelism shards embedding and LM-head rows over TP ranks.
+    // Hidden states remain replicated; embedding outputs and CE hidden
+    // gradients are summed over the TP communicator.
+    bool vocab_parallel = false;
+    int64_t local_vocab_size = 0;
+    // Set when a legacy NCCL setter supplies an incompatible mixed topology.
+    // Training entry points reject the context before touching parameters.
+    bool topology_invalid = false;
+    // A failed transactional recovery quarantines this context. Continuing
+    // would risk rank-divergent NCCL collectives or tenant state corruption.
+    bool poisoned = false;
+    int cuda_device = 0;
 // ──────────────────────────────────────────────────────────────────────
 };
 
-// ── Multi-LoRA: concat all adapters' A/B, 2x GEMM ──
-// Pre-build cache of concatenated A/B per (layer, module) pair.
-// Called once at start of forward; reused across all layers.
+static bool find_canonical_adapter_index(
+    const TrainingContext* ctx, int64_t adapter_id, size_t& adapter_index
+) {
+    if (!ctx || !ctx->adapter_id_index_valid ||
+            ctx->adapter_id_index.size() != ctx->adapters.size())
+        return false;
+    const auto it = std::lower_bound(
+        ctx->adapter_id_index.begin(), ctx->adapter_id_index.end(), adapter_id,
+        [](const auto& entry, int64_t id) { return entry.first < id; });
+    if (it == ctx->adapter_id_index.end() || it->first != adapter_id ||
+            it->second >= ctx->adapters.size() ||
+            ctx->adapters[it->second].id != adapter_id)
+        return false;
+    adapter_index = it->second;
+    return true;
+}
 
-static void precompute_lora_cache(TrainingContext* ctx) {
-    if (ctx->lora_cache_valid) return;
-    ctx->lora_cache.clear();
+static void require_canonical_adapter_index(const TrainingContext* ctx) {
+    TORCH_CHECK(ctx && ctx->adapter_id_index_valid &&
+            ctx->adapter_id_index.size() == ctx->adapters.size(),
+        "dynamic LoRA adapter ID index is unavailable or inconsistent");
+}
 
-    // Phase 1: collect all (cache_key, a_concat, b_concat) tuples
-    struct LoraEntry {
-        int64_t key;
-        at::Tensor a_concat;  // [sum_ranks, in]
-        at::Tensor b_concat;  // [out, sum_ranks]
-    };
-    std::vector<LoraEntry> entries;
+static cudaStream_t context_moe_shared_stream(TrainingContext* ctx) {
+    TORCH_CHECK(ctx, "shared-expert overlap requires a training context");
+    if (ctx->moe_shared_stream) return ctx->moe_shared_stream;
+    int previous_device = 0;
+    TORCH_CHECK(cudaGetDevice(&previous_device) == cudaSuccess,
+        "failed to query CUDA device for shared-expert stream");
+    if (previous_device != ctx->cuda_device) {
+        TORCH_CHECK(cudaSetDevice(ctx->cuda_device) == cudaSuccess,
+            "failed to select CUDA device for shared-expert stream");
+    }
+    const auto create_status = cudaStreamCreateWithFlags(
+        &ctx->moe_shared_stream, cudaStreamNonBlocking);
+    if (previous_device != ctx->cuda_device) cudaSetDevice(previous_device);
+    TORCH_CHECK(create_status == cudaSuccess,
+        "failed to create context-owned shared-expert stream");
+    return ctx->moe_shared_stream;
+}
 
-    for (int64_t layer_idx = 0; layer_idx < ctx->num_layers; layer_idx++) {
-        int64_t num_pairs = (ctx->layer_configs[layer_idx].layer_type == 0) ? 4 : 3;
-        for (int64_t pair_idx = 0; pair_idx < num_pairs; pair_idx++) {
-            std::vector<at::Tensor> a_list, b_list;
-            for (auto& adapter : ctx->adapters) {
-                if (!adapter.target_layers.empty() && adapter.target_layers.find(layer_idx) == adapter.target_layers.end())
-                    continue;
-                auto it = adapter.params.find(layer_idx);
-                if (it == adapter.params.end()) continue;
-                if (pair_idx >= (int64_t)it->second.size()) continue;
-                auto& [a, b] = it->second[pair_idx];
-                double scaling = adapter.alpha / (double)adapter.rank;
-                b_list.push_back(b * scaling);
-                a_list.push_back(a);
-            }
-            if (a_list.empty()) {
-                if (!ctx->lora_a.empty() && layer_idx < (int64_t)ctx->lora_layer_offset.size()) {
-                    int64_t la_offset = ctx->lora_layer_offset[layer_idx];
-                    if (la_offset + pair_idx < (int64_t)ctx->lora_a.size()) {
-                        b_list.push_back(ctx->lora_b[la_offset + pair_idx] * ctx->lora_scaling);
-                        a_list.push_back(ctx->lora_a[la_offset + pair_idx]);
-                    }
-                }
-            }
-            if (!a_list.empty()) {
-                at::Tensor a_concat, b_concat;
-                if (a_list.size() == 1) {
-                    a_concat = a_list[0];
-                    b_concat = b_list[0];
-                } else {
-                    a_concat = at::cat(a_list, 0);  // [sum_ranks, in]
-                    b_concat = at::cat(b_list, 1);  // [out, sum_ranks]
-                }
-                entries.push_back({layer_idx * 10 + pair_idx, a_concat, b_concat});
-            }
-        }
+static cudaStream_t context_moe_ep_metadata_stream(TrainingContext* ctx) {
+    TORCH_CHECK(ctx, "EP count overlap requires a training context");
+    if (ctx->moe_ep_metadata_stream) return ctx->moe_ep_metadata_stream;
+    int previous_device = 0;
+    TORCH_CHECK(cudaGetDevice(&previous_device) == cudaSuccess,
+        "failed to query CUDA device for EP metadata stream");
+    if (previous_device != ctx->cuda_device) {
+        TORCH_CHECK(cudaSetDevice(ctx->cuda_device) == cudaSuccess,
+            "failed to select CUDA device for EP metadata stream");
     }
+    const auto create_status = cudaStreamCreateWithFlags(
+        &ctx->moe_ep_metadata_stream, cudaStreamNonBlocking);
+    if (previous_device != ctx->cuda_device) cudaSetDevice(previous_device);
+    TORCH_CHECK(create_status == cudaSuccess,
+        "failed to create context-owned EP metadata stream");
+    return ctx->moe_ep_metadata_stream;
+}
 
-    // Phase 2: group by (out_dim, in_dim, sum_ranks) and batch matmul
-    // delta = b_concat @ a_concat  → [out, in]
-    // Group entries with identical shapes to use at::bmm
-    struct ShapeGroup {
-        int64_t out_dim, in_dim, sum_ranks;
-        std::vector<size_t> indices;
+static at::Tensor derive_attention_mask(
+    TrainingContext* ctx, const at::Tensor& input_ids
+) {
+    TORCH_CHECK(ctx && ctx->pad_token_id >= 0,
+        "attention mask is omitted but no pad_token_id is configured");
+    TORCH_CHECK(input_ids.defined() && input_ids.dim() == 2 &&
+            input_ids.scalar_type() == at::kLong,
+        "cannot derive attention mask from invalid input_ids");
+    return input_ids.ne(ctx->pad_token_id).to(at::kFloat).contiguous();
+}
+
+static at::Tensor fused_full_attention_qkv_weight(
+    TrainingContext* ctx, int64_t layer_idx
+) {
+    if (!ctx || layer_idx < 0 || layer_idx >= static_cast<int64_t>(
+            ctx->fused_full_attention_qkv_weights.size()))
+        return at::Tensor();
+    return ctx->fused_full_attention_qkv_weights[layer_idx];
+}
+
+static at::Tensor fused_mlp_fc1_weight(
+    TrainingContext* ctx, int64_t layer_idx
+) {
+    if (!ctx || layer_idx < 0 || layer_idx >= static_cast<int64_t>(
+            ctx->fused_mlp_fc1_weights.size()))
+        return at::Tensor();
+    return ctx->fused_mlp_fc1_weights[layer_idx];
+}
+
+struct RouterAuxLossFunction
+    : public torch::autograd::Function<RouterAuxLossFunction> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* autograd_ctx,
+        at::Tensor routing_weights,
+        at::Tensor aux_loss,
+        at::Tensor backward_scale
+    ) {
+        autograd_ctx->save_for_backward({aux_loss, backward_scale});
+        return routing_weights;
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* autograd_ctx,
+        std::vector<at::Tensor> grad_outputs
+    ) {
+        auto saved = autograd_ctx->get_saved_variables();
+        return {
+            grad_outputs[0],
+            at::ones_like(saved[0]) * saved[1],
+            at::Tensor(),
+        };
+    }
+};
+
+static void router_aux_allreduce(
+    TrainingContext* ctx, at::Tensor& tensor, const char* operation
+) {
+    auto stream = c10::cuda::getCurrentCUDAStream(
+        tensor.device().index()).stream();
+    auto reduce = [&](ncclComm_t comm, int world, const char* axis) {
+        if (!comm || world <= 1) return;
+        auto err = ncclAllReduce(
+            tensor.data_ptr<float>(), tensor.data_ptr<float>(), tensor.numel(),
+            ncclFloat, ncclSum, comm, stream);
+        TORCH_CHECK(err == ncclSuccess, operation, " ", axis,
+            " all-reduce failed: ", ncclGetErrorString(err));
     };
-    std::vector<ShapeGroup> groups;
-    for (size_t i = 0; i < entries.size(); i++) {
-        auto& e = entries[i];
-        int64_t out_dim = e.b_concat.size(0);
-        int64_t sum_ranks = e.b_concat.size(1);
-        int64_t in_dim = e.a_concat.size(1);
-        bool found = false;
-        for (auto& g : groups) {
-            if (g.out_dim == out_dim && g.in_dim == in_dim && g.sum_ranks == sum_ranks) {
-                g.indices.push_back(i);
-                found = true;
-                break;
-            }
-        }
-        if (!found) {
-            groups.push_back({out_dim, in_dim, sum_ranks, {i}});
-        }
+    // EP and DP own distinct token shards. TP keeps hidden states replicated,
+    // so reducing over TP would count the same routing decisions twice.
+    reduce(ctx->nccl_comm, ctx->ep_world_size, "EP");
+    reduce(ctx->dp_comm, ctx->dp_world_size, "DP");
+}
+
+static at::Tensor attach_router_aux_loss(
+    TrainingContext* ctx,
+    const at::Tensor& routing_weights,
+    const at::Tensor& topk_indices,
+    int64_t batch,
+    int64_t seq,
+    int64_t top_k,
+    int64_t num_experts
+) {
+    if (!ctx || ctx->router_aux_loss_coef == 0.0 ||
+            (!ctx->collect_router_aux_loss && !at::GradMode::is_enabled()))
+        return routing_weights;
+
+    TORCH_CHECK(ctx->adapters.empty(),
+        "router auxiliary loss is not yet supported for dynamic multi-LoRA; "
+        "per-tenant routing statistics are required to preserve isolation");
+    TORCH_CHECK(!ctx->has_mtp,
+        "router auxiliary loss with MTP is not yet supported");
+    TORCH_CHECK(ctx->pp_world_size == 1,
+        "router auxiliary loss with pipeline parallelism is not yet supported");
+    TORCH_CHECK(routing_weights.dim() == 2 &&
+            routing_weights.size(0) == batch * seq &&
+            routing_weights.size(1) == num_experts,
+        "router auxiliary loss received invalid probability shape");
+
+    at::Tensor valid;
+    if (ctx->attention_mask.defined() && ctx->attention_mask.numel() > 0) {
+        TORCH_CHECK(ctx->attention_mask.dim() == 2 &&
+                ctx->attention_mask.size(0) == batch &&
+                ctx->attention_mask.size(1) == seq,
+            "router auxiliary loss attention-mask shape mismatch");
+        valid = ctx->attention_mask.reshape({batch * seq}).ne(0);
+    } else {
+        valid = at::ones(
+            {batch * seq}, routing_weights.options().dtype(at::kBool));
     }
 
-    // Phase 3: for each group, stack and bmm in BF16 (2x faster tensor cores)
-    // LoRA params are FP32 for gradient stability; cast to BF16 for matmul only.
-    // Autograd handles the cast backward (grad → FP32 automatically).
-    auto bf16 = at::kBFloat16;
-    for (auto& g : groups) {
-        int n = (int)g.indices.size();
-        if (n == 1) {
-            auto& e = entries[g.indices[0]];
-            auto delta = at::matmul(e.b_concat, e.a_concat);
-            ctx->lora_cache[e.key] = delta;  // already BF16
-        } else {
-            std::vector<at::Tensor> b_stack_vec, a_stack_vec;
-            b_stack_vec.reserve(n);
-            a_stack_vec.reserve(n);
-            for (auto idx : g.indices) {
-                b_stack_vec.push_back(entries[idx].b_concat);
-                a_stack_vec.push_back(entries[idx].a_concat);
-            }
-            auto b_stack = at::stack(b_stack_vec, 0);  // [N, out, sum_ranks] BF16
-            auto a_stack = at::stack(a_stack_vec, 0);  // [N, sum_ranks, in] BF16
-            auto deltas = at::bmm(b_stack, a_stack);   // [N, out, in] BF16
-            for (int i = 0; i < n; i++) {
-                ctx->lora_cache[entries[g.indices[i]].key] = deltas[i];
-            }
-        }
+    auto selected = topk_indices.masked_select(valid.unsqueeze(-1));
+    auto tokens_per_expert = at::bincount(
+        selected, c10::nullopt, num_experts).to(at::kFloat);
+    router_aux_allreduce(ctx, tokens_per_expert, "router token-count");
+
+    auto valid_probs = routing_weights *
+        valid.to(routing_weights.scalar_type()).unsqueeze(-1);
+    auto total_tokens =
+        (tokens_per_expert.sum() / static_cast<double>(top_k)).clamp_min(1.0);
+    auto aux_loss =
+        (valid_probs.sum(0) * tokens_per_expert).sum() *
+        (static_cast<double>(num_experts) * ctx->router_aux_loss_coef /
+            static_cast<double>(top_k)) /
+        (total_tokens * total_tokens);
+
+    if (ctx->collect_router_aux_loss) {
+        auto detached = aux_loss.detach();
+        ctx->router_aux_loss_value = ctx->router_aux_loss_value.defined()
+            ? ctx->router_aux_loss_value + detached
+            : detached;
     }
+    return RouterAuxLossFunction::apply(
+        routing_weights, aux_loss, ctx->router_aux_backward_scale);
+}
 
-    ctx->lora_cache_valid = true;
+static void begin_router_aux_loss(
+    TrainingContext* ctx, double backward_scale
+) {
+    if (!ctx || ctx->router_aux_loss_coef == 0.0) return;
+    TORCH_CHECK(std::isfinite(backward_scale) && backward_scale >= 0.0,
+        "router auxiliary backward scale must be finite and non-negative");
+    ctx->router_aux_backward_scale = at::full(
+        {}, backward_scale,
+        at::TensorOptions().device(at::kCUDA, ctx->cuda_device).dtype(at::kFloat));
+    router_aux_allreduce(
+        ctx, ctx->router_aux_backward_scale, "router backward-scale");
+    ctx->router_aux_loss_value = at::Tensor();
+    ctx->collect_router_aux_loss = true;
 }
 
-// ── Batched Multi-LoRA: activation-level B@(A@x) ──
+static double finish_router_aux_loss(TrainingContext* ctx) {
+    if (!ctx || ctx->router_aux_loss_coef == 0.0) return 0.0;
+    ctx->collect_router_aux_loss = false;
+    if (!ctx->router_aux_loss_value.defined()) return 0.0;
+    auto global_value = ctx->router_aux_loss_value.detach().clone();
+    router_aux_allreduce(ctx, global_value, "router loss-value");
+    const double value = global_value.item<double>();
+    ctx->router_aux_loss_value = at::Tensor();
+    return value;
+}
 
-/// Prepare stacked A/B tensors for all adapters per (layer, module).
-/// Stores in ctx->lora_batch_cache. Called once before forward.
-/// Replaces precompute_lora_cache when N > 1.
-static void prepare_lora_batch(TrainingContext* ctx) {
-    if (ctx->lora_batch_valid) return;
-    ctx->lora_batch_cache.clear();
-    ctx->lora_batch_n = 0;
+struct AdapterRegistryHash {
+    uint64_t first = 1469598103934665603ULL;
+    uint64_t second = 1099511628211ULL;
 
-    for (int64_t layer_idx = 0; layer_idx < ctx->num_layers; layer_idx++) {
-        int64_t num_pairs = (ctx->layer_configs[layer_idx].layer_type == 0) ? 4 : 3;
-        for (int64_t pair_idx = 0; pair_idx < num_pairs; pair_idx++) {
-            std::vector<at::Tensor> a_list, b_list;
-            std::vector<double> scalings;
-            for (auto& adapter : ctx->adapters) {
-                if (!adapter.target_layers.empty() && adapter.target_layers.find(layer_idx) == adapter.target_layers.end())
-                    continue;
-                auto it = adapter.params.find(layer_idx);
-                if (it == adapter.params.end()) continue;
-                if (pair_idx >= (int64_t)it->second.size()) continue;
-                auto& [a, b] = it->second[pair_idx];
-                b_list.push_back(b);
-                a_list.push_back(a);
-                scalings.push_back(adapter.alpha / (double)adapter.rank);
-            }
-            if (a_list.empty()) continue;
+    void add_u64(uint64_t value) {
+        first ^= value;
+        first *= 1099511628211ULL;
+        second ^= value + 0x9e3779b97f4a7c15ULL + (second << 6) +
+            (second >> 2);
+        second *= 0xbf58476d1ce4e5b9ULL;
+    }
 
-            int64_t n = (int64_t)a_list.size();
-            if (ctx->lora_batch_n == 0) ctx->lora_batch_n = n;
+    void add_string(const std::string& value) {
+        add_u64(value.size());
+        for (const unsigned char byte : value) add_u64(byte);
+    }
+};
 
-            auto a_stack = at::stack(a_list, 0);  // [N, rank, in]
-            auto b_stack = at::stack(b_list, 0);  // [N, out, rank]
-            // Create scaling tensor on GPU — from_blob only wraps CPU pointer,
-            // so we must explicitly move it to the right device.
-            auto scaling_cpu = at::from_blob(
-                scalings.data(), {(int64_t)n, 1, 1},
-                at::TensorOptions().dtype(at::kDouble)
-            );
-            auto scaling = scaling_cpu.to(a_stack.device()).to(at::kBFloat16);  // [N, 1, 1]
+static void hash_double(AdapterRegistryHash& hash, double value) {
+    uint64_t bits = 0;
+    static_assert(sizeof(bits) == sizeof(value));
+    std::memcpy(&bits, &value, sizeof(bits));
+    hash.add_u64(bits);
+}
 
-            ctx->lora_batch_cache[layer_idx * 10 + pair_idx] = {
-                a_stack, b_stack, scaling
-            };
-        }
+
+static bool global_to_local_layer(
+    const TrainingContext* ctx, int64_t global_layer, int64_t& local_layer
+) {
+    if (!ctx || global_layer < ctx->global_layer_start ||
+            global_layer >= ctx->global_layer_start + ctx->num_layers) {
+        return false;
     }
-    ctx->lora_batch_valid = true;
+    local_layer = global_layer - ctx->global_layer_start;
+    return true;
 }
 
-/// Compute B@(A@x) * scaling — never materializes B@A.
-/// x: [N, seq, in], A: [N, rank, in], B: [N, out, rank], scaling: [N, 1, 1]
-/// returns: [N, seq, out]
-static at::Tensor lora_activation_delta(
-    const at::Tensor& x,          // [N, seq, in]
-    const at::Tensor& A,          // [N, rank, in]
-    const at::Tensor& B,          // [N, out, rank]
-    const at::Tensor& scaling     // [N, 1, 1]
+static void hash_tensor_layout(
+    AdapterRegistryHash& hash, const at::Tensor& tensor
 ) {
-    // Cast to compute dtype (BF16)
-    auto kind = x.scalar_type();
-    auto A_c = A.to(kind);
-    auto B_c = B.to(kind);
-    auto s_c = scaling.to(kind);
-    // Ax = A @ x^T  → [N, rank, seq]
-    auto Ax = at::bmm(A_c, x.transpose(-2, -1));
-    // delta = B @ Ax → [N, out, seq] → transpose → [N, seq, out]
-    auto delta = at::bmm(B_c, Ax).transpose(-2, -1);
-    return delta * s_c;
+    hash.add_u64(tensor.defined());
+    if (!tensor.defined()) return;
+    hash.add_u64(static_cast<uint64_t>(tensor.scalar_type()));
+    hash.add_u64(tensor.requires_grad());
+    hash.add_u64(tensor.dim());
+    for (const auto size : tensor.sizes()) hash.add_u64(size);
 }
 
-__attribute__((noinline, visibility("default")))
-at::Tensor apply_multi_lora(
-    TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx,
-    const at::Tensor& base_weight
+static void hash_adapter_layout(
+    AdapterRegistryHash& hash,
+    const TrainingContext::LoRAAdapter& adapter
 ) {
-    auto it = ctx->lora_cache.find(layer_idx * 10 + pair_idx);
-    if (it == ctx->lora_cache.end()) return base_weight;
+    hash.add_u64(adapter.id);
+    hash.add_u64(adapter.rank);
+    hash.add_u64(adapter.optimizer_step);
+    hash_double(hash, adapter.optimizer_lr);
+    hash_double(hash, adapter.optimizer_beta1);
+    hash_double(hash, adapter.optimizer_beta2);
+    hash_double(hash, adapter.optimizer_eps);
+    hash_double(hash, adapter.alpha);
+    hash.add_u64(adapter.all_target_layers);
+    hash.add_u64(adapter.global_target_layers.size());
+    for (const auto layer : adapter.global_target_layers) hash.add_u64(layer);
+    hash.add_u64(adapter.target_modules.size());
+    for (const auto& module : adapter.target_modules) hash.add_string(module);
+    hash.add_u64(adapter.params.size());
+    for (const auto& [layer, pairs] : adapter.params) {
+        hash.add_u64(layer);
+        hash.add_u64(pairs.size());
+        for (const auto& [a, b] : pairs) {
+            hash_tensor_layout(hash, a);
+            hash_tensor_layout(hash, b);
+        }
+    }
+    hash.add_u64(adapter.grad_accum.size());
+    for (const auto& [layer, pairs] : adapter.grad_accum) {
+        hash.add_u64(layer);
+        hash.add_u64(pairs.size());
+        for (const auto& pair : pairs) {
+            hash_tensor_layout(hash, pair[0]);
+            hash_tensor_layout(hash, pair[1]);
+        }
+    }
+    hash_tensor_layout(hash, adapter.grad_slab);
+}
 
-    // Cached delta_weight = b_concat @ a_concat (BF16, precomputed in batched bmm)
-    auto& delta = it->second;
-    return base_weight + delta;  // both BF16, no conversion needed
+static void hash_adapter_request_identity(
+    AdapterRegistryHash& hash,
+    const TrainingContext::LoRAAdapter& adapter
+) {
+    hash.add_u64(adapter.id);
+    hash.add_u64(adapter.rank);
+    hash.add_u64(adapter.optimizer_step);
+    hash_double(hash, adapter.optimizer_lr);
+    hash_double(hash, adapter.optimizer_beta1);
+    hash_double(hash, adapter.optimizer_beta2);
+    hash_double(hash, adapter.optimizer_eps);
+    hash_double(hash, adapter.alpha);
+    hash.add_u64(adapter.all_target_layers);
+    hash.add_u64(adapter.global_target_layers.size());
+    for (const auto layer : adapter.global_target_layers) hash.add_u64(layer);
+    hash.add_u64(adapter.target_modules.size());
+    for (const auto& module : adapter.target_modules) hash.add_string(module);
 }
 
-// Forward declarations for sub-layer checkpointing
-// Sub-layer checkpointing: split each layer into attn + mlp segments
-// Enabled by QWEN36_SUBCKPT=1 env var. Reduces peak memory by ~2x
-// at the cost of 2x extra recomputation per layer during backward.
-// ──────────────────────────────────────────────────────────────────────
+static void hash_collective_runtime_environment(AdapterRegistryHash& hash) {
+    hash.add_u64(env_enabled("QWEN36_EP_A2A"));
+    hash.add_u64(env_enabled("QWEN36_EP_A2A_SHARDED"));
+    hash.add_u64(env_enabled("QWEN36_EP_A2A_PACKED", true));
+    hash.add_u64(env_enabled("QWEN36_EP_A2A_PACKED_METADATA", true));
+    hash.add_u64(env_enabled("QWEN36_EP_A2A_COUNT_OVERLAP"));
+    hash.add_u64(env_enabled("QWEN36_MOE_FUSED_UNPERMUTE"));
+    hash.add_u64(env_enabled("QWEN36_MOE_FUSED_LOCAL_PERMUTE"));
+    hash.add_u64(env_enabled("QWEN36_DISABLE_GROUPED_MM"));
+    hash.add_u64(env_enabled("QWEN36_GROUPED_LORA_SYNC", true));
+    hash.add_u64(env_enabled("QWEN36_PACKED_LORA_SYNC", true));
+    hash.add_u64(env_enabled("QWEN36_GRAD_SLAB", true));
+    hash.add_u64(env_enabled(
+        "QWEN36_LAZY_DYNAMIC_OPTIMIZER_STATE", true));
+    hash.add_u64(env_enabled(
+        "QWEN36_DYNAMIC_ADAM_SHADOW_POOL", true));
+    hash.add_u64(env_enabled("QWEN36_HETERO_PADDED_BATCH", true));
+    hash.add_u64(env_enabled("QWEN36_GDN_RECURRENT_FUSION", true));
+    hash.add_u64(env_enabled("QWEN36_GDN_CHUNKWISE_BWD"));
+    hash.add_u64(env_enabled("QWEN36_GDN_INVERSE_BWD"));
+    hash.add_u64(env_enabled("QWEN36_GDN_FUSED_AB_PROJECTION", true));
+    hash.add_u64(env_enabled("QWEN36_GDN_FUSED_CP_EXCHANGE"));
+    hash.add_u64(env_enabled("QWEN36_DELTA_REFERENCE_BWD"));
+    hash.add_u64(env_enabled("QWEN36_SUBCKPT"));
+    hash.add_u64(env_enabled("QWEN36_FUSED_LAYER"));
+    hash.add_u64(env_enabled("QWEN36_FUSED_CE"));
+    hash.add_u64(env_enabled("QWEN36_FUSED_QKV"));
+    hash.add_u64(env_enabled("QWEN36_FUSED_LORA_QKV_A"));
+    hash.add_u64(env_enabled("QWEN36_FUSED_MLP_FC1"));
+    hash.add_u64(env_enabled("QWEN36_MOE_SHARED_OVERLAP"));
+    hash.add_u64(env_enabled("QWEN36_CP_FULL_ATTENTION_KV_GATHER"));
+    hash.add_u64(env_enabled("QWEN36_CP_FULL_ATTENTION_RING"));
+    hash.add_u64(env_enabled("QWEN36_DISABLE_MTP"));
+    const char* checkpoint_stride =
+        getenv("QWEN36_GDN_STATE_CHECKPOINT_STRIDE");
+    hash.add_string(checkpoint_stride ? checkpoint_stride : "");
+    const char* sequence_chunk = getenv("QWEN36_SEQ_CHUNK");
+    hash.add_string(sequence_chunk ? sequence_chunk : "");
+    const char* ce_token_tile = getenv("QWEN36_CE_TOKEN_TILE");
+    hash.add_string(ce_token_tile ? ce_token_tile : "");
+    const char* checkpoint_group = getenv("QWEN36_GROUP_SIZE");
+    hash.add_string(checkpoint_group ? checkpoint_group : "");
+}
 
-// Compute attention output from hidden
-at::Tensor compute_attn_only(
-    TrainingContext* ctx, const at::Tensor& hidden, int64_t layer_idx, at::ScalarType kind
+static void hash_collective_runtime_config(
+    AdapterRegistryHash& hash, const TrainingContext* ctx
 ) {
-    const auto& cfg = ctx->layer_configs[layer_idx];
-    int64_t w_offset = 0;
-    for (int64_t j = 0; j < layer_idx; j++)
-        w_offset += weight_count_for_layer(ctx->layer_configs[j]);
-    auto attn_input = rms_norm(hidden, *ctx->weight_ptrs[w_offset + 0], cfg.rms_eps);
-
-    // Use batched path if lora_batch is active
-    if (ctx->lora_batch_valid) {
-        if (cfg.layer_type == 0) {
-            auto qp = *ctx->weight_ptrs[w_offset+2], qn = *ctx->weight_ptrs[w_offset+3];
-            auto kp = *ctx->weight_ptrs[w_offset+4], kn = *ctx->weight_ptrs[w_offset+5];
-            auto vp = *ctx->weight_ptrs[w_offset+6], op = *ctx->weight_ptrs[w_offset+7];
-            return full_attention_batched(
-                ctx, attn_input, layer_idx, qp, qn, kp, kn, vp, op,
-                cfg.num_heads, cfg.num_kv_heads, cfg.head_dim,
-                cfg.partial_rotary_factor, cfg.rope_theta, cfg.rms_eps, kind,
-                ctx->attention_mask);
-        } else {
-            auto qkv = *ctx->weight_ptrs[w_offset+2], z = *ctx->weight_ptrs[w_offset+3];
-            auto a = *ctx->weight_ptrs[w_offset+4], b = *ctx->weight_ptrs[w_offset+5];
-            auto al = *ctx->weight_ptrs[w_offset+6], db = *ctx->weight_ptrs[w_offset+7];
-            auto cw = *ctx->weight_ptrs[w_offset+8], nw = *ctx->weight_ptrs[w_offset+9];
-            auto op = *ctx->weight_ptrs[w_offset+10];
-            return linear_attention_batched(
-                ctx, attn_input, layer_idx, qkv, z, a, b, al, db, cw, nw, op,
-                cfg.num_k_heads, cfg.key_dim, cfg.num_v_heads, cfg.val_dim,
-                cfg.conv_kernel, cfg.rms_eps, kind);
+    hash_collective_runtime_environment(hash);
+    hash.add_u64(ctx->use_checkpoint);
+    hash.add_u64(ctx->group_size);
+    hash_double(hash, ctx->max_grad_norm);
+    hash.add_u64(ctx->has_mtp);
+    hash_double(hash, ctx->mtp_loss_scale);
+    if (ctx->has_mtp) {
+        hash.add_u64(ctx->vocab_size);
+        hash.add_u64(ctx->local_vocab_size);
+        for (const auto* tensor : {
+                ctx->embed_ptr.empty() ? nullptr : ctx->embed_ptr[0],
+                ctx->lm_head_ptr.empty() ? nullptr : ctx->lm_head_ptr[0]}) {
+            hash.add_u64(tensor != nullptr);
+            if (tensor) hash_tensor_layout(hash, *tensor);
         }
-    }
-
-    // Legacy path: weight-level LoRA
-    int64_t lora_count = (cfg.layer_type == 0) ? 4 : 3;
-    int64_t la_offset = ctx->lora_layer_offset[layer_idx];
-    bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size();
-    std::vector<at::Tensor*> la(lora_count, nullptr), lb(lora_count, nullptr);
-    if (has_lora) for (int64_t k = 0; k < lora_count; k++) { la[k] = &ctx->lora_a[la_offset + k]; lb[k] = &ctx->lora_b[la_offset + k]; }
-
-    if (cfg.layer_type == 0) {
-        auto qp = *ctx->weight_ptrs[w_offset+2], qn = *ctx->weight_ptrs[w_offset+3];
-        auto kp = *ctx->weight_ptrs[w_offset+4], kn = *ctx->weight_ptrs[w_offset+5];
-        auto vp = *ctx->weight_ptrs[w_offset+6], op = *ctx->weight_ptrs[w_offset+7];
-        if (has_lora) {
-            if (la[0]) qp = lora_delta(qp, *la[0], *lb[0], ctx->lora_scaling);
-            if (la[1]) kp = lora_delta(kp, *la[1], *lb[1], ctx->lora_scaling);
-            if (la[2]) vp = lora_delta(vp, *la[2], *lb[2], ctx->lora_scaling);
-            if (la[3]) op = lora_delta(op, *la[3], *lb[3], ctx->lora_scaling);
+        for (const auto* tensor : {
+                ctx->mtp_fc, ctx->mtp_pre_fc_norm_emb,
+                ctx->mtp_pre_fc_norm_hidden, ctx->mtp_norm}) {
+            hash.add_u64(tensor != nullptr);
+            if (tensor) hash_tensor_layout(hash, *tensor);
         }
-        return full_attention(attn_input, qp, qn, kp, kn, vp, op,
-            cfg.num_heads, cfg.num_kv_heads, cfg.head_dim,
-            cfg.partial_rotary_factor, cfg.rope_theta, cfg.rms_eps, kind,
-            ctx->attention_mask);
-    } else {
-        auto qkv = *ctx->weight_ptrs[w_offset+2], z = *ctx->weight_ptrs[w_offset+3];
-        auto a = *ctx->weight_ptrs[w_offset+4], b = *ctx->weight_ptrs[w_offset+5];
-        auto al = *ctx->weight_ptrs[w_offset+6], db = *ctx->weight_ptrs[w_offset+7];
-        auto cw = *ctx->weight_ptrs[w_offset+8], nw = *ctx->weight_ptrs[w_offset+9];
-        auto op = *ctx->weight_ptrs[w_offset+10];
-        if (has_lora) {
-            if (la[0]) qkv = lora_delta(qkv, *la[0], *lb[0], ctx->lora_scaling);
-            if (la[1]) z = lora_delta(z, *la[1], *lb[1], ctx->lora_scaling);
-            if (la[2]) op = lora_delta(op, *la[2], *lb[2], ctx->lora_scaling);
+        hash.add_u64(ctx->mtp_layer_configs.size());
+        for (const auto& config : ctx->mtp_layer_configs) {
+            hash.add_u64(config.layer_type);
+            hash.add_u64(config.num_heads);
+            hash.add_u64(config.num_kv_heads);
+            hash.add_u64(config.head_dim);
+            hash.add_u64(config.num_k_heads);
+            hash.add_u64(config.key_dim);
+            hash.add_u64(config.num_v_heads);
+            hash.add_u64(config.val_dim);
+            hash.add_u64(config.conv_kernel);
+            hash_double(hash, config.partial_rotary_factor);
+            hash_double(hash, config.rope_theta);
+            hash_double(hash, config.rms_eps);
+            hash.add_u64(config.num_experts);
+            hash.add_u64(config.top_k);
+            hash.add_u64(config.moe_intermediate);
+            hash.add_u64(config.expert_count);
+            const bool ownership_valid = config.num_experts <= 0 ||
+                (ctx->ep_world_size > 0 &&
+                    config.num_experts % ctx->ep_world_size == 0 &&
+                    config.expert_count ==
+                        config.num_experts / ctx->ep_world_size &&
+                    config.expert_start ==
+                        ctx->ep_rank * config.expert_count);
+            hash.add_u64(ownership_valid);
+            hash.add_u64(config.intermediate_size);
+            hash.add_u64(config.norm_topk_prob);
+        }
+        hash.add_u64(ctx->mtp_layer_weights.size());
+        for (const auto* tensor : ctx->mtp_layer_weights) {
+            hash.add_u64(tensor != nullptr);
+            if (tensor) hash_tensor_layout(hash, *tensor);
         }
-        return linear_attention(attn_input, qkv, z, a, b, al, db, cw, nw, op,
-            cfg.num_k_heads, cfg.key_dim, cfg.num_v_heads, cfg.val_dim,
-            cfg.conv_kernel, cfg.rms_eps, kind);
     }
 }
 
-// Compute MLP output from residual (hidden + attn_output)
-at::Tensor compute_mlp_only(
-    TrainingContext* ctx, const at::Tensor& residual, int64_t layer_idx, at::ScalarType kind
+static void hash_collective_topology(
+    AdapterRegistryHash& hash, const TrainingContext* ctx
 ) {
-    const auto& cfg = ctx->layer_configs[layer_idx];
-    int64_t w_offset = 0;
-    for (int64_t j = 0; j < layer_idx; j++)
-        w_offset += weight_count_for_layer(ctx->layer_configs[j]);
-    int64_t mlp_start = (cfg.layer_type == 0) ? 8 : 11;
-    auto post_attn = rms_norm(residual, *ctx->weight_ptrs[w_offset + 1], cfg.rms_eps);
-    if (cfg.num_experts > 0) {
-        return moe_forward(cfg.nccl_comm, cfg.nccl_stream, post_attn,
-            *ctx->weight_ptrs[w_offset+mlp_start], *ctx->weight_ptrs[w_offset+mlp_start+1],
-            *ctx->weight_ptrs[w_offset+mlp_start+2], *ctx->weight_ptrs[w_offset+mlp_start+3],
-            *ctx->weight_ptrs[w_offset+mlp_start+4], *ctx->weight_ptrs[w_offset+mlp_start+5],
-            *ctx->weight_ptrs[w_offset+mlp_start+6],
-            cfg.num_experts, cfg.top_k, cfg.moe_intermediate,
-            cfg.norm_topk_prob != 0, cfg.expert_start, cfg.expert_count, kind);
-    } else {
-        return dense_mlp_forward(post_attn,
-            *ctx->weight_ptrs[w_offset+mlp_start], *ctx->weight_ptrs[w_offset+mlp_start+1],
-            *ctx->weight_ptrs[w_offset+mlp_start+2], kind);
+    hash.add_u64(ctx->num_layers);
+    hash.add_u64(ctx->global_layer_start);
+    hash.add_u64(ctx->global_num_layers);
+    hash.add_u64(ctx->tp_world_size);
+    hash.add_u64(ctx->cp_world_size);
+    hash.add_u64(ctx->ep_world_size);
+    hash.add_u64(ctx->dp_world_size);
+    hash.add_u64(ctx->pp_world_size);
+    hash.add_u64(ctx->base_tp_attention);
+    hash.add_u64(ctx->base_tp_mlp);
+    hash.add_u64(ctx->vocab_parallel);
+    hash.add_u64(ctx->sequence_parallel);
+    hash.add_u64(ctx->expert_parallel);
+    hash.add_u64(ctx->data_parallel);
+    hash_collective_runtime_config(hash, ctx);
+    if (ctx->cp_world_size > 1) {
+        for (const auto& config : ctx->layer_configs) {
+            hash.add_u64(config.layer_type);
+            hash.add_u64(config.num_heads);
+            hash.add_u64(config.num_kv_heads);
+            hash.add_u64(config.head_dim);
+            hash.add_u64(config.num_k_heads);
+            hash.add_u64(config.key_dim);
+            hash.add_u64(config.num_v_heads);
+            hash.add_u64(config.val_dim);
+            hash.add_u64(config.conv_kernel);
+            hash_double(hash, config.partial_rotary_factor);
+            hash_double(hash, config.rope_theta);
+            hash_double(hash, config.rms_eps);
+            hash.add_u64(config.num_experts);
+            hash.add_u64(config.top_k);
+            hash.add_u64(config.moe_intermediate);
+            hash.add_u64(config.expert_start);
+            hash.add_u64(config.expert_count);
+            hash.add_u64(config.intermediate_size);
+            hash.add_u64(config.norm_topk_prob);
+        }
+        hash.add_u64(ctx->weight_ptrs.size());
+        for (const auto* weight : ctx->weight_ptrs) {
+            hash.add_u64(weight != nullptr);
+            if (weight) hash_tensor_layout(hash, *weight);
+        }
+        for (const auto* tensor : {
+                ctx->embed_ptr.empty() ? nullptr : ctx->embed_ptr[0],
+                ctx->final_norm_ptr.empty() ? nullptr : ctx->final_norm_ptr[0],
+                ctx->lm_head_ptr.empty() ? nullptr : ctx->lm_head_ptr[0]}) {
+            hash.add_u64(tensor != nullptr);
+            if (tensor) hash_tensor_layout(hash, *tensor);
+        }
     }
+    uint64_t router_aux_bits = 0;
+    static_assert(sizeof(router_aux_bits) == sizeof(ctx->router_aux_loss_coef));
+    std::memcpy(
+        &router_aux_bits, &ctx->router_aux_loss_coef,
+        sizeof(router_aux_bits));
+    hash.add_u64(router_aux_bits);
 }
 
-// Sub-layer checkpoint: wraps a function call with no-grad forward + recompute on backward
-struct SubLayerCkpt : public torch::autograd::Function<SubLayerCkpt> {
-    static at::Tensor forward(
-        torch::autograd::AutogradContext* ctx,
-        at::Tensor input, int64_t tc_val, int64_t layer_idx, bool is_attn
-    ) {
-        ctx->saved_data["tc"] = tc_val;
-        ctx->saved_data["layer"] = layer_idx;
-        ctx->saved_data["is_attn"] = is_attn;
-        bool offload = getenv("QWEN36_OFFLOAD_ACTIVATIONS");
-        if (offload) {
-            ctx->saved_data["input_cpu"] = input.detach().to(
-                at::TensorOptions().dtype(input.scalar_type()).device(at::kCPU).pinned_memory(true));
-            ctx->saved_data["device"] = input.device();
-        } else {
-            ctx->save_for_backward({input});
+static void validate_native_execution_topology(
+    const TrainingContext* ctx, bool allow_pipeline = false) {
+    TORCH_CHECK(ctx, "native Qwen execution requires a valid context");
+    TORCH_CHECK(gdn_state_checkpoint_environment_valid(),
+        "QWEN36_GDN_STATE_CHECKPOINT_STRIDE=0 is only valid with "
+        "QWEN36_GDN_INVERSE_BWD or QWEN36_DELTA_REFERENCE_BWD and not "
+        "QWEN36_GDN_CHUNKWISE_BWD; otherwise use unset or an integer >= 2");
+    TORCH_CHECK(!ctx->topology_invalid,
+        "native Qwen context rejected an incompatible TP/CP/EP/DP/PP topology");
+    TORCH_CHECK(ctx->tp_world_size == 1 || ctx->tp_comm,
+        "native Qwen TP communicator is not initialized");
+    TORCH_CHECK(ctx->cp_world_size == 1 || ctx->cp_comm,
+        "native Qwen CP communicator is not initialized");
+    TORCH_CHECK(ctx->ep_world_size == 1 || ctx->nccl_comm,
+        "native Qwen EP communicator is not initialized");
+    TORCH_CHECK(ctx->dp_world_size == 1 || ctx->dp_comm,
+        "native Qwen DP communicator is not initialized");
+    TORCH_CHECK(ctx->pp_world_size == 1 || ctx->pp_comm,
+        "native Qwen PP communicator is not initialized");
+    TORCH_CHECK(allow_pipeline || ctx->pp_world_size == 1,
+        "native Qwen pipeline execution must use the pipeline training ABI");
+    if (ctx->has_mtp) {
+        if (ctx->tp_world_size > 1) {
+            TORCH_CHECK(ctx->tp_world_size == 2 && ctx->tp_comm &&
+                    ctx->base_tp_attention && ctx->base_tp_mlp,
+                "MTP tensor parallelism currently requires initialized "
+                "TP_SIZE=2 with both base attention and MLP tensor parallelism");
         }
-        at::AutoGradMode guard(false);
-        auto* tc = reinterpret_cast<TrainingContext*>(tc_val);
-        if (is_attn) {
-            // Attn segment: returns hidden + attn_output (residual connection)
-            return input + compute_attn_only(tc, input, layer_idx, tc->compute_type);
-        } else {
-            // MLP segment: returns residual + mlp_out (full layer output)
-            return input + compute_mlp_only(tc, input, layer_idx, tc->compute_type);
+        TORCH_CHECK(!ctx->sequence_parallel &&
+                ctx->cp_world_size == 1 && ctx->pp_world_size == 1,
+            "MTP parallelism currently requires replicated sequence and "
+            "vocabulary-compatible state with CP=PP=1");
+        if (ctx->ep_world_size > 1) {
+            TORCH_CHECK(ctx->expert_parallel && ctx->nccl_comm &&
+                    env_enabled("QWEN36_EP_A2A") &&
+                    env_enabled("QWEN36_EP_A2A_SHARDED") &&
+                    env_enabled("QWEN36_EP_A2A_PACKED", true),
+                "MTP expert parallelism requires initialized packed sharded "
+                "A2A");
+        }
+        for (const auto& config : ctx->mtp_layer_configs) {
+            TORCH_CHECK(config.layer_type == 0,
+                "MTP parallelism currently supports full-attention "
+                "prediction layers only");
+            TORCH_CHECK(ctx->ep_world_size == 1 || config.num_experts > 0,
+                "MTP expert parallelism requires MoE prediction layers");
         }
     }
-    static std::vector<at::Tensor> backward(
-        torch::autograd::AutogradContext* ctx, std::vector<at::Tensor> grad_output
-    ) {
-        at::Tensor input;
-        if (ctx->saved_data.count("input_cpu") > 0) {
-            input = ctx->saved_data["input_cpu"].toTensor().to(ctx->saved_data["device"].toDevice());
-        } else {
-            input = ctx->get_saved_variables()[0];
+    if (ctx->sequence_parallel) {
+        TORCH_CHECK(ctx->tp_world_size == 2 && ctx->tp_comm,
+            "sequence parallelism currently requires initialized TP_SIZE=2");
+        TORCH_CHECK(ctx->cp_world_size == 1 && ctx->pp_world_size == 1 &&
+                ctx->ep_world_size == 1 && ctx->dp_world_size == 1,
+            "sequence parallelism currently requires CP=PP=EP=DP=1");
+        TORCH_CHECK(ctx->base_tp_attention && ctx->base_tp_mlp &&
+                ctx->vocab_parallel,
+            "sequence parallelism requires base attention, base MLP, and "
+            "vocabulary tensor parallelism");
+        TORCH_CHECK(!ctx->has_mtp,
+            "sequence parallelism does not support MTP");
+        TORCH_CHECK(ctx->router_aux_loss_coef == 0.0,
+            "sequence parallelism does not support router auxiliary loss");
+        for (const auto& config : ctx->layer_configs) {
+            TORCH_CHECK(config.num_experts == 0,
+                "sequence parallelism currently supports dense layers only");
         }
-        auto* tc = reinterpret_cast<TrainingContext*>(ctx->saved_data["tc"].toInt());
-        int64_t layer = ctx->saved_data["layer"].toInt();
-        bool is_attn = ctx->saved_data["is_attn"].toBool();
-        at::AutoGradMode guard(true);
-        input.set_requires_grad(true);
-        auto output = is_attn ? compute_attn_only(tc, input, layer, tc->compute_type)
-                              : compute_mlp_only(tc, input, layer, tc->compute_type);
-
-        // Collect all tensors to compute gradients for: input + LoRA params for this layer
-        // This way grad() accumulates gradients into LoRA params (leaf nodes) too.
-        int64_t lora_count = (tc->layer_configs[layer].layer_type == 0) ? 4 : 3;
-        int64_t la_offset = tc->lora_layer_offset[layer];
-        bool has_lora = (la_offset + lora_count) <= (int64_t)tc->lora_a.size();
-
-        std::vector<at::Tensor> grad_inputs = {input};
-        if (has_lora) {
-            for (int64_t k = 0; k < lora_count; k++) {
-                grad_inputs.push_back(tc->lora_a[la_offset + k]);
-                grad_inputs.push_back(tc->lora_b[la_offset + k]);
+        for (int64_t layer = 0; layer < ctx->num_layers; ++layer) {
+            const auto table = lora_projection_table(ctx->layer_configs[layer]);
+            const int64_t offset = ctx->lora_layer_offset[layer];
+            for (int64_t pair = 0; pair < table.count; ++pair) {
+                const int64_t slot = offset + pair;
+                if (slot >= static_cast<int64_t>(ctx->lora_active.size()) ||
+                    !ctx->lora_active[slot]) continue;
+                TORCH_CHECK(
+                    lora_tp_layout(ctx, layer, pair) != LoraTpLayout::LatentRank,
+                    "sequence parallelism does not support latent-rank LoRA; "
+                    "projection ", table.entries[pair].name,
+                    " at layer ", layer, " must use base TP sharding");
             }
         }
-
-        auto grads = torch::autograd::grad(
-            {output}, grad_inputs, {grad_output[0]},
-            /*retain_graph=*/false, /*create_graph=*/false,
-            /*allow_unused=*/true
-        );
-
-        // Manually accumulate LoRA param gradients
-        if (has_lora) {
-            int64_t gi = 1;  // skip input grad (index 0)
-            for (int64_t k = 0; k < lora_count; k++) {
-                if (grads[gi].defined()) {
-                    auto& param_a = tc->lora_a[la_offset + k];
-                    if (param_a.grad().defined())
-                        param_a.grad().add_(grads[gi]);
-                    else
-                        param_a.mutable_grad() = grads[gi].clone();
-                }
-                gi++;
-                if (grads[gi].defined()) {
-                    auto& param_b = tc->lora_b[la_offset + k];
-                    if (param_b.grad().defined())
-                        param_b.grad().add_(grads[gi]);
-                    else
-                        param_b.mutable_grad() = grads[gi].clone();
+        // Dynamic adapters use the same projection-aware TP geometry as the
+        // fixed registry. Sequence-parallel activation sharding cannot split
+        // latent-rank factors because those factors are owned by TP ranks.
+        for (const auto& adapter : ctx->adapters) {
+            for (const auto& [layer, pairs] : adapter.params) {
+                TORCH_CHECK(layer >= 0 &&
+                        layer < static_cast<int64_t>(ctx->layer_configs.size()),
+                    "dynamic LoRA adapter ", adapter.id,
+                    " references an invalid layer: ", layer);
+                const auto table = lora_projection_table(ctx->layer_configs[layer]);
+                TORCH_CHECK(pairs.size() <= table.entries.size(),
+                    "dynamic LoRA adapter ", adapter.id,
+                    " has an invalid projection registry at layer ", layer);
+                for (int64_t pair = 0;
+                     pair < static_cast<int64_t>(pairs.size()); ++pair) {
+                    if (!pairs[pair].first.defined() ||
+                            !pairs[pair].first.requires_grad()) continue;
+                    TORCH_CHECK(
+                        lora_tp_layout(ctx, layer, pair) != LoraTpLayout::LatentRank,
+                        "sequence parallelism does not support latent-rank dynamic LoRA; "
+                        "adapter ", adapter.id, " projection ",
+                        table.entries[pair].name,
+                        " at layer ", layer,
+                        " must use base TP sharding");
                 }
-                gi++;
             }
         }
-
-        return {grads[0], at::Tensor(), at::Tensor(), at::Tensor()};
     }
-};
+    if (ctx->cp_world_size > 1) {
+        TORCH_CHECK(ctx->cp_comm,
+            "Qwen context parallelism requires an initialized CP communicator");
+        TORCH_CHECK(ctx->tp_world_size == 1 && ctx->ep_world_size == 1 &&
+                ctx->dp_world_size == 1 && ctx->pp_world_size == 1,
+            "Qwen context parallelism currently requires TP=EP=DP=PP=1");
+        TORCH_CHECK(!ctx->sequence_parallel && !ctx->base_tp_attention &&
+                !ctx->base_tp_mlp && !ctx->vocab_parallel,
+            "Qwen context parallelism cannot be combined with TP or sequence parallelism");
+        TORCH_CHECK(!ctx->has_mtp,
+            "Qwen context parallelism does not support MTP");
+        TORCH_CHECK(ctx->adapters.empty() ||
+                env_enabled("QWEN36_GROUPED_LORA_SYNC", true),
+            "dynamic Qwen context parallelism requires grouped LoRA gradient sync");
+        TORCH_CHECK(ctx->router_aux_loss_coef == 0.0,
+            "Qwen context parallelism does not support router auxiliary loss");
+        const char* sequence_chunk = getenv("QWEN36_SEQ_CHUNK");
+        TORCH_CHECK(!sequence_chunk || atoll(sequence_chunk) <= 0,
+            "Qwen context parallelism does not support QWEN36_SEQ_CHUNK");
+        int64_t weight_offset = 0;
+        for (int64_t layer = 0; layer < ctx->num_layers; ++layer) {
+            const auto& config = ctx->layer_configs[layer];
+            TORCH_CHECK(config.num_experts == 0,
+                "Qwen context parallelism currently requires dense layers");
+            if (config.layer_type == 0) {
+                const bool use_kv_gather = env_enabled(
+                    "QWEN36_CP_FULL_ATTENTION_KV_GATHER");
+                const bool use_ring = env_enabled(
+                    "QWEN36_CP_FULL_ATTENTION_RING");
+                TORCH_CHECK(use_kv_gather != use_ring,
+                    "full-attention context parallelism requires exactly one of "
+                    "QWEN36_CP_FULL_ATTENTION_KV_GATHER=1 or "
+                    "QWEN36_CP_FULL_ATTENTION_RING=1");
+                TORCH_CHECK(ctx->cp_world_size == 2 || use_ring,
+                    "CP_SIZE>2 full attention requires "
+                    "QWEN36_CP_FULL_ATTENTION_RING=1");
+                TORCH_CHECK(config.num_heads > 0 &&
+                        config.num_kv_heads > 0 && config.head_dim > 0 &&
+                        config.num_heads % config.num_kv_heads == 0 &&
+                        config.partial_rotary_factor >= 0.0 &&
+                        config.partial_rotary_factor <= 1.0,
+                    "full-attention CP head geometry is invalid");
+                const int64_t rotary_dim = static_cast<int64_t>(
+                    config.head_dim * config.partial_rotary_factor);
+                TORCH_CHECK(rotary_dim >= 0 &&
+                        rotary_dim <= config.head_dim &&
+                        rotary_dim % 2 == 0,
+                    "full-attention CP rotary dimension must be even and "
+                    "no larger than head_dim");
+                auto* q = ctx->weight_ptrs[weight_offset + 2];
+                auto* q_norm = ctx->weight_ptrs[weight_offset + 3];
+                auto* k = ctx->weight_ptrs[weight_offset + 4];
+                auto* k_norm = ctx->weight_ptrs[weight_offset + 5];
+                auto* v = ctx->weight_ptrs[weight_offset + 6];
+                auto* out = ctx->weight_ptrs[weight_offset + 7];
+                TORCH_CHECK(q && q_norm && k && k_norm && v && out &&
+                        q->dim() == 2 && k->dim() == 2 && v->dim() == 2 &&
+                        out->dim() == 2 && q_norm->dim() == 1 &&
+                        k_norm->dim() == 1,
+                    "full-attention CP received missing or invalid weights");
+                const int64_t hidden = q->size(1);
+                const int64_t q_size =
+                    config.num_heads * config.head_dim;
+                const int64_t kv_size =
+                    config.num_kv_heads * config.head_dim;
+                TORCH_CHECK(q->sizes() == at::IntArrayRef({
+                            q_size * 2, hidden}) &&
+                        q_norm->size(0) == config.head_dim &&
+                        k->sizes() == at::IntArrayRef({kv_size, hidden}) &&
+                        k_norm->size(0) == config.head_dim &&
+                        v->sizes() == at::IntArrayRef({kv_size, hidden}) &&
+                        out->sizes() == at::IntArrayRef({hidden, q_size}),
+                    "full-attention CP weight shapes do not match the model "
+                    "configuration");
+                weight_offset += weight_count_for_layer(config);
+                continue;
+            }
+            TORCH_CHECK(ctx->cp_world_size == 2,
+                "CP_SIZE>2 currently supports full-attention-only models; "
+                "hybrid GDN context parallelism remains restricted to CP2");
+            TORCH_CHECK(config.num_k_heads > 0 && config.num_v_heads > 0 &&
+                    config.key_dim > 0 && config.val_dim > 0 &&
+                    config.conv_kernel > 0 &&
+                    config.num_v_heads % config.num_k_heads == 0 &&
+                    config.num_k_heads % ctx->cp_world_size == 0 &&
+                    config.num_v_heads % ctx->cp_world_size == 0,
+                "GDN head counts must preserve groups and be divisible by CP_SIZE");
+            auto* qkv = ctx->weight_ptrs[weight_offset + 2];
+            auto* z = ctx->weight_ptrs[weight_offset + 3];
+            auto* a = ctx->weight_ptrs[weight_offset + 4];
+            auto* b = ctx->weight_ptrs[weight_offset + 5];
+            auto* a_log = ctx->weight_ptrs[weight_offset + 6];
+            auto* dt_bias = ctx->weight_ptrs[weight_offset + 7];
+            auto* conv = ctx->weight_ptrs[weight_offset + 8];
+            auto* norm = ctx->weight_ptrs[weight_offset + 9];
+            auto* out = ctx->weight_ptrs[weight_offset + 10];
+            TORCH_CHECK(qkv && z && a && b && a_log && dt_bias && conv &&
+                    norm && out && qkv->dim() == 2,
+                "GDN context parallelism received missing or invalid weights");
+            const int64_t q = config.num_k_heads * config.key_dim;
+            const int64_t v = config.num_v_heads * config.val_dim;
+            const int64_t hidden = qkv->size(1);
+            TORCH_CHECK(qkv->size(0) == q * 2 + v &&
+                    z->dim() == 2 && z->sizes() == at::IntArrayRef({v, hidden}) &&
+                    a->dim() == 2 &&
+                    a->sizes() == at::IntArrayRef({config.num_v_heads, hidden}) &&
+                    b->sizes() == a->sizes() &&
+                    a_log->dim() == 1 &&
+                    a_log->size(0) == config.num_v_heads &&
+                    dt_bias->sizes() == a_log->sizes() &&
+                    conv->dim() == 3 && conv->size(0) == q * 2 + v &&
+                    conv->size(1) == 1 &&
+                    conv->size(2) == config.conv_kernel &&
+                    norm->dim() == 1 && norm->size(0) == config.val_dim &&
+                    out->dim() == 2 &&
+                    out->sizes() == at::IntArrayRef({hidden, v}),
+                "GDN context-parallel weight shapes do not match the model configuration");
+            weight_offset += weight_count_for_layer(config);
+        }
+    }
+}
 
-// Forward a single layer with sub-layer checkpointing
-// attn segment returns (hidden + attn_output) to avoid extra GPU tensor
-// mlp segment returns (residual + mlp_out) = full layer output
-at::Tensor forward_single_layer_subckpt(
-    TrainingContext* ctx, const at::Tensor& hidden, int64_t layer_idx
+// Adapter requests are canonical across pipeline stages.  A PP stage owns a
+// different contiguous layer slice, so the local layer count/start and local
+// tensor layouts must not participate in the request identity.  The full
+// topology hash above remains the stricter check for ranks that share tensor
+// ownership (TP/EP/DP).
+static void hash_collective_request_topology(
+    AdapterRegistryHash& hash, const TrainingContext* ctx
 ) {
-    // attn segment: computes attn_output, returns hidden + attn_output
-    auto residual = SubLayerCkpt::apply(
-        hidden, (int64_t)(uintptr_t)ctx, layer_idx, true);
-    // mlp segment: computes mlp_out from residual, returns residual + mlp_out
-    auto result = SubLayerCkpt::apply(
-        residual, (int64_t)(uintptr_t)ctx, layer_idx, false);
-    return result;
+    hash.add_u64(ctx->global_num_layers);
+    hash.add_u64(ctx->tp_world_size);
+    hash.add_u64(ctx->cp_world_size);
+    hash.add_u64(ctx->ep_world_size);
+    hash.add_u64(ctx->dp_world_size);
+    hash.add_u64(ctx->pp_world_size);
+    hash.add_u64(ctx->base_tp_attention);
+    hash.add_u64(ctx->base_tp_mlp);
+    hash.add_u64(ctx->vocab_parallel);
+    hash.add_u64(ctx->sequence_parallel);
+    hash.add_u64(ctx->expert_parallel);
+    hash.add_u64(ctx->data_parallel);
+    hash_collective_runtime_config(hash, ctx);
 }
 
-// ── Batched attention variants (activation-level LoRA) ──
-// These wrap the base attention functions, applying LoRA delta as
-// B@(A@x) * scaling on the projected outputs instead of modifying weights.
-
-static at::Tensor full_attention_batched(
-    TrainingContext* ctx, const at::Tensor& hidden, int64_t layer_idx,
-    const at::Tensor& q_proj, const at::Tensor& q_norm,
-    const at::Tensor& k_proj, const at::Tensor& k_norm,
-    const at::Tensor& v_proj, const at::Tensor& o_proj,
-    int64_t num_heads, int64_t num_kv_heads, int64_t head_dim,
-    double partial_rotary_factor, double rope_theta,
-    double rms_eps, at::ScalarType kind,
-    const at::Tensor& attention_mask
+static void validate_adapter_collective_registry(
+    TrainingContext* ctx,
+    const int64_t* requested_ids,
+    int64_t requested_count,
+    int64_t requested_rank,
+    bool use_registered_order,
+    const int64_t* expected_steps = nullptr
 ) {
-    // Compute Q/K/V with base weight, then add LoRA delta if present
-    int64_t batch = hidden.size(0), seq = hidden.size(1);
-    int64_t qkv_dim = num_heads * head_dim;
+    if (!ctx || (!ctx->nccl_comm && !ctx->dp_comm && !ctx->tp_comm &&
+            !ctx->cp_comm)) return;
+
+    AdapterRegistryHash hash;
+    hash_collective_topology(hash, ctx);
+    hash.add_u64(requested_count);
+    hash.add_u64(requested_rank);
+    hash.add_u64(use_registered_order);
+    hash.add_u64(expected_steps != nullptr);
+    int64_t found_count = 0;
+    if (use_registered_order) {
+        hash.add_u64(ctx->adapters.size());
+        for (const auto& adapter : ctx->adapters) {
+            hash_adapter_layout(hash, adapter);
+            ++found_count;
+        }
+    } else if (requested_ids && requested_count > 0) {
+        for (int64_t index = 0; index < requested_count; ++index) {
+            const int64_t requested_id = requested_ids[index];
+            hash.add_u64(requested_id);
+            if (expected_steps) hash.add_u64(expected_steps[index]);
+            size_t adapter_index = 0;
+            if (!find_canonical_adapter_index(
+                    ctx, requested_id, adapter_index)) {
+                hash.add_u64(0x6d697373696e67ULL);
+                continue;
+            }
+            hash.add_u64(0x666f756e64ULL);
+            hash_adapter_layout(hash, ctx->adapters[adapter_index]);
+            ++found_count;
+        }
+    } else {
+        hash.add_u64(0x6e756c6cULL);
+    }
 
-    auto q = at::matmul(hidden, q_proj.t());
-    auto k = at::matmul(hidden, k_proj.t());
-    auto v = at::matmul(hidden, v_proj.t());
+    constexpr uint64_t kPositiveInt64Mask =
+        static_cast<uint64_t>(std::numeric_limits<int64_t>::max());
+    const std::vector<int64_t> signature_values{
+        requested_count,
+        requested_rank,
+        found_count,
+        static_cast<int64_t>(hash.first & kPositiveInt64Mask),
+        static_cast<int64_t>(hash.second & kPositiveInt64Mask),
+    };
+    c10::cuda::set_device(ctx->cuda_device);
+    cudaSetDevice(ctx->cuda_device);
+    auto options = at::TensorOptions().dtype(at::kLong).device(
+        at::kCUDA, ctx->cuda_device);
+    auto minimum = at::tensor(signature_values, options);
+    auto maximum = minimum.clone();
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+
+    auto reduce_axis = [&](ncclComm_t communicator, const char* axis) {
+        if (!communicator) return;
+        auto err = ncclAllReduce(
+            minimum.data_ptr<int64_t>(), minimum.data_ptr<int64_t>(),
+            minimum.numel(), ncclInt64, ncclMin, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " adapter registry minimum all-reduce failed: ",
+            ncclGetErrorString(err));
+        err = ncclAllReduce(
+            maximum.data_ptr<int64_t>(), maximum.data_ptr<int64_t>(),
+            maximum.numel(), ncclInt64, ncclMax, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " adapter registry maximum all-reduce failed: ",
+            ncclGetErrorString(err));
+    };
 
-    // Apply activation-level LoRA: q += B@(A@hidden) * scaling
-    auto it_q = ctx->lora_batch_cache.find(layer_idx * 10 + 0);
-    if (it_q != ctx->lora_batch_cache.end()) {
-        q = q + lora_activation_delta(hidden, it_q->second.a_stack, it_q->second.b_stack, it_q->second.scaling);
+    // Sequential reductions over the orthogonal axes propagate the extrema
+    // over the complete TP x EP x DP grid.
+    reduce_axis(ctx->nccl_comm, "EP");
+    reduce_axis(ctx->dp_comm, "DP");
+    reduce_axis(ctx->tp_comm, "TP");
+    reduce_axis(ctx->cp_comm, "CP");
+    const auto minimum_cpu = minimum.to(at::kCPU);
+    const auto maximum_cpu = maximum.to(at::kCPU);
+    const auto* minimum_data = minimum_cpu.data_ptr<int64_t>();
+    const auto* maximum_data = maximum_cpu.data_ptr<int64_t>();
+    for (int64_t index = 0; index < minimum.numel(); ++index) {
+        TORCH_CHECK(minimum_data[index] == maximum_data[index],
+            "dynamic LoRA adapter registry mismatch across distributed ranks; "
+            "all ranks must select the same ordered IDs, ranks, optimizer "
+            "clocks, targets, and tensor layouts");
     }
-    auto it_k = ctx->lora_batch_cache.find(layer_idx * 10 + 1);
-    if (it_k != ctx->lora_batch_cache.end()) {
-        k = k + lora_activation_delta(hidden, it_k->second.a_stack, it_k->second.b_stack, it_k->second.scaling);
+}
+
+static void validate_fixed_collective_registry(
+    TrainingContext* ctx, int64_t optimizer_phase = -1
+) {
+    if (!ctx || (!ctx->nccl_comm && !ctx->dp_comm && !ctx->tp_comm &&
+            !ctx->cp_comm)) return;
+
+    AdapterRegistryHash hash;
+    hash_collective_topology(hash, ctx);
+    hash.add_u64(ctx->fixed_optimizer_step);
+    hash.add_u64(ctx->lora_active.size());
+    hash.add_u64(ctx->lora_a.size());
+    hash.add_u64(ctx->lora_b.size());
+    hash.add_u64(ctx->adapters.size());
+    for (const auto& adapter : ctx->adapters)
+        hash_adapter_layout(hash, adapter);
+    for (size_t index = 0; index < ctx->lora_active.size(); ++index) {
+        hash.add_u64(ctx->lora_active[index]);
+        if (index < ctx->lora_a.size())
+            hash_tensor_layout(hash, ctx->lora_a[index]);
+        if (index < ctx->lora_b.size())
+            hash_tensor_layout(hash, ctx->lora_b[index]);
+        if (index < ctx->grad_accum_a.size())
+            hash_tensor_layout(hash, ctx->grad_accum_a[index]);
+        if (index < ctx->grad_accum_b.size())
+            hash_tensor_layout(hash, ctx->grad_accum_b[index]);
     }
-    auto it_v = ctx->lora_batch_cache.find(layer_idx * 10 + 2);
-    if (it_v != ctx->lora_batch_cache.end()) {
-        v = v + lora_activation_delta(hidden, it_v->second.a_stack, it_v->second.b_stack, it_v->second.scaling);
+    hash_tensor_layout(hash, ctx->fixed_grad_slab);
+
+    constexpr uint64_t kPositiveInt64Mask =
+        static_cast<uint64_t>(std::numeric_limits<int64_t>::max());
+    const std::vector<int64_t> signature_values{
+        optimizer_phase,
+        ctx->fixed_optimizer_step,
+        static_cast<int64_t>(ctx->lora_active.size()),
+        static_cast<int64_t>(hash.first & kPositiveInt64Mask),
+        static_cast<int64_t>(hash.second & kPositiveInt64Mask),
+    };
+    c10::cuda::set_device(ctx->cuda_device);
+    cudaSetDevice(ctx->cuda_device);
+    auto options = at::TensorOptions().dtype(at::kLong).device(
+        at::kCUDA, ctx->cuda_device);
+    auto minimum = at::tensor(signature_values, options);
+    auto maximum = minimum.clone();
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+    auto reduce_axis = [&](ncclComm_t communicator, const char* axis) {
+        if (!communicator) return;
+        auto err = ncclAllReduce(
+            minimum.data_ptr<int64_t>(), minimum.data_ptr<int64_t>(),
+            minimum.numel(), ncclInt64, ncclMin, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " fixed LoRA registry minimum all-reduce failed: ",
+            ncclGetErrorString(err));
+        err = ncclAllReduce(
+            maximum.data_ptr<int64_t>(), maximum.data_ptr<int64_t>(),
+            maximum.numel(), ncclInt64, ncclMax, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " fixed LoRA registry maximum all-reduce failed: ",
+            ncclGetErrorString(err));
+    };
+    reduce_axis(ctx->nccl_comm, "EP");
+    reduce_axis(ctx->dp_comm, "DP");
+    reduce_axis(ctx->tp_comm, "TP");
+    reduce_axis(ctx->cp_comm, "CP");
+
+    const auto minimum_cpu = minimum.to(at::kCPU);
+    const auto maximum_cpu = maximum.to(at::kCPU);
+    const auto* minimum_data = minimum_cpu.data_ptr<int64_t>();
+    const auto* maximum_data = maximum_cpu.data_ptr<int64_t>();
+    for (int64_t index = 0; index < minimum.numel(); ++index) {
+        TORCH_CHECK(minimum_data[index] == maximum_data[index],
+            "fixed LoRA registry mismatch across distributed ranks; all "
+            "ranks must use the same active slots, tensor layouts, and "
+            "optimizer phase/clock");
     }
+}
 
-    // Reshape Q: [batch, seq, num_heads, head_dim*2] → split into q and gate
-    q = q.view({batch, seq, num_heads, head_dim * 2});
-    auto qk = q.chunk(2, -1);
-    auto q_out = qk[0].transpose(1, 2);   // [batch, heads, seq, head_dim]
-    auto gate = qk[1].transpose(1, 2);
-
-    k = k.view({batch, seq, num_kv_heads, head_dim}).transpose(1, 2);
-    v = v.view({batch, seq, num_kv_heads, head_dim}).transpose(1, 2);
-
-    q_out = rms_norm(q_out, q_norm, rms_eps);
-    k = rms_norm(k, k_norm, rms_eps);
+static void validate_pipeline_collective_registry(
+    TrainingContext* ctx,
+    int64_t optimizer_phase,
+    double micro_token_weight,
+    double next_accumulated_token_weight
+) {
+    if (!ctx || !ctx->pp_comm || ctx->pp_world_size <= 1) return;
+
+    AdapterRegistryHash hash;
+    hash.add_u64(ctx->global_num_layers);
+    hash.add_u64(ctx->fixed_all_target_layers);
+    hash.add_u64(ctx->fixed_target_layers_global.size());
+    for (const auto layer : ctx->fixed_target_layers_global)
+        hash.add_u64(layer);
+    hash.add_u64(ctx->fixed_target_modules.size());
+    for (const auto& module : ctx->fixed_target_modules)
+        hash.add_string(module);
+    hash_collective_runtime_config(hash, ctx);
+    hash.add_u64(ctx->fixed_optimizer_step);
+    uint64_t token_weight_bits = 0;
+    static_assert(sizeof(token_weight_bits) == sizeof(ctx->accumulated_token_weight));
+    std::memcpy(&token_weight_bits, &ctx->accumulated_token_weight,
+        sizeof(token_weight_bits));
+    hash.add_u64(token_weight_bits);
+    std::memcpy(&token_weight_bits, &micro_token_weight,
+        sizeof(token_weight_bits));
+    hash.add_u64(token_weight_bits);
+    std::memcpy(&token_weight_bits, &next_accumulated_token_weight,
+        sizeof(token_weight_bits));
+    hash.add_u64(token_weight_bits);
+
+    constexpr uint64_t kPositiveInt64Mask =
+        static_cast<uint64_t>(std::numeric_limits<int64_t>::max());
+    const std::vector<int64_t> signature_values{
+        optimizer_phase,
+        ctx->fixed_optimizer_step,
+        static_cast<int64_t>(hash.first & kPositiveInt64Mask),
+        static_cast<int64_t>(hash.second & kPositiveInt64Mask),
+    };
+    auto options = at::TensorOptions().dtype(at::kLong).device(
+        at::kCUDA, ctx->cuda_device);
+    auto minimum = at::tensor(signature_values, options);
+    auto maximum = minimum.clone();
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+    auto control_comm = ctx->pp_control_comm ?
+        ctx->pp_control_comm : ctx->pp_comm;
+    auto min_error = ncclAllReduce(
+        minimum.data_ptr<int64_t>(), minimum.data_ptr<int64_t>(),
+        minimum.numel(), ncclInt64, ncclMin, control_comm, stream);
+    TORCH_CHECK(min_error == ncclSuccess,
+        "PP fixed-LoRA registry minimum all-reduce failed: ",
+        ncclGetErrorString(min_error));
+    auto max_error = ncclAllReduce(
+        maximum.data_ptr<int64_t>(), maximum.data_ptr<int64_t>(),
+        maximum.numel(), ncclInt64, ncclMax, control_comm, stream);
+    TORCH_CHECK(max_error == ncclSuccess,
+        "PP fixed-LoRA registry maximum all-reduce failed: ",
+        ncclGetErrorString(max_error));
+    const auto minimum_cpu = minimum.to(at::kCPU);
+    const auto maximum_cpu = maximum.to(at::kCPU);
+    const auto* minimum_data = minimum_cpu.data_ptr<int64_t>();
+    const auto* maximum_data = maximum_cpu.data_ptr<int64_t>();
+    for (int64_t index = 0; index < minimum.numel(); ++index) {
+        TORCH_CHECK(minimum_data[index] == maximum_data[index],
+            "pipeline fixed-LoRA registry mismatch across stages; all stages "
+            "must use the same global target identity, runtime configuration, "
+            "optimizer phase, clock, microstep token weight, and accumulation window");
+    }
+}
 
-    // Release gate before RoPE
-    gate = at::Tensor();
+static void validate_pipeline_window_collective_registry(
+    TrainingContext* ctx,
+    int64_t window_id,
+    int64_t num_microbatches,
+    int32_t schedule,
+    int32_t num_chunks,
+    int32_t flags
+) {
+    TORCH_CHECK(ctx && (ctx->pp_control_comm || ctx->pp_comm) &&
+            ctx->pp_world_size >= 2,
+        "pipeline window requires a PP communicator with PP_SIZE >= 2");
+    AdapterRegistryHash hash;
+    hash.add_u64(ctx->global_num_layers);
+    if (flags == kPipelineWindowFlagDynamicLora) {
+        hash.add_u64(ctx->adapters.size());
+        for (const auto& adapter : ctx->adapters)
+            hash_adapter_request_identity(hash, adapter);
+        hash.add_u64(ctx->pad_heterogeneous_lora_batch);
+    }
+    hash.add_u64(ctx->fixed_all_target_layers);
+    hash.add_u64(ctx->fixed_target_layers_global.size());
+    for (const auto layer : ctx->fixed_target_layers_global)
+        hash.add_u64(layer);
+    hash.add_u64(ctx->fixed_target_modules.size());
+    for (const auto& module : ctx->fixed_target_modules)
+        hash.add_string(module);
+    hash_collective_runtime_config(hash, ctx);
+    hash.add_u64(ctx->fixed_optimizer_step);
+    uint64_t token_weight_bits = 0;
+    std::memcpy(&token_weight_bits, &ctx->accumulated_token_weight,
+        sizeof(token_weight_bits));
+    hash.add_u64(token_weight_bits);
+    hash.add_u64(static_cast<uint64_t>(window_id));
+    hash.add_u64(static_cast<uint64_t>(num_microbatches));
+    hash.add_u64(static_cast<uint64_t>(schedule));
+    hash.add_u64(static_cast<uint64_t>(num_chunks));
+    hash.add_u64(static_cast<uint64_t>(flags));
+
+    constexpr uint64_t kPositiveInt64Mask =
+        static_cast<uint64_t>(std::numeric_limits<int64_t>::max());
+    const std::vector<int64_t> values{
+        window_id,
+        num_microbatches,
+        schedule,
+        num_chunks,
+        flags,
+        ctx->fixed_optimizer_step,
+        static_cast<int64_t>(hash.first & kPositiveInt64Mask),
+        static_cast<int64_t>(hash.second & kPositiveInt64Mask),
+    };
+    auto options = at::TensorOptions().dtype(at::kLong).device(
+        at::kCUDA, ctx->cuda_device);
+    auto minimum = at::tensor(values, options);
+    auto maximum = minimum.clone();
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+    auto control_comm = ctx->pp_control_comm ?
+        ctx->pp_control_comm : ctx->pp_comm;
+    auto min_error = ncclAllReduce(
+        minimum.data_ptr<int64_t>(), minimum.data_ptr<int64_t>(),
+        minimum.numel(), ncclInt64, ncclMin, control_comm, stream);
+    TORCH_CHECK(min_error == ncclSuccess,
+        "pipeline window registry minimum all-reduce failed: ",
+        ncclGetErrorString(min_error));
+    auto max_error = ncclAllReduce(
+        maximum.data_ptr<int64_t>(), maximum.data_ptr<int64_t>(),
+        maximum.numel(), ncclInt64, ncclMax, control_comm, stream);
+    TORCH_CHECK(max_error == ncclSuccess,
+        "pipeline window registry maximum all-reduce failed: ",
+        ncclGetErrorString(max_error));
+    const auto minimum_cpu = minimum.to(at::kCPU);
+    const auto maximum_cpu = maximum.to(at::kCPU);
+    const auto* minimum_data = minimum_cpu.data_ptr<int64_t>();
+    const auto* maximum_data = maximum_cpu.data_ptr<int64_t>();
+    for (int64_t index = 0; index < minimum.numel(); ++index) {
+        TORCH_CHECK(minimum_data[index] == maximum_data[index],
+            "pipeline window registry mismatch across stages; all stages "
+            "must use the same runtime configuration and execute the same "
+            "window and schedule");
+    }
+}
 
-    // RoPE
-    int64_t rotary_dim = (int64_t)(head_dim * partial_rotary_factor);
-    if (rotary_dim > 0) {
-        auto device = hidden.device();
-        auto pos = at::arange(seq, at::TensorOptions().dtype(at::kFloat).device(device)).unsqueeze(0);
-        auto exponents = at::arange(0, rotary_dim, 2, at::TensorOptions().dtype(at::kFloat).device(device)) / (double)rotary_dim;
-        auto inv_freq = (exponents * std::log(rope_theta)).exp().reciprocal();
-        auto freqs = pos.unsqueeze(-1) * inv_freq.unsqueeze(0);
-        auto emb = at::cat({freqs, freqs}, -1);
-        auto cos = emb.cos().unsqueeze(1).to(q_out.scalar_type());
-        auto sin = emb.sin().unsqueeze(1).to(q_out.scalar_type());
-        auto q_rot = q_out.narrow(-1, 0, rotary_dim);
-        auto k_rot = k.narrow(-1, 0, rotary_dim);
-        auto rotate_half_q = at::cat({-q_rot.narrow(-1, rotary_dim/2, rotary_dim/2), q_rot.narrow(-1, 0, rotary_dim/2)}, -1);
-        auto rotate_half_k = at::cat({-k_rot.narrow(-1, rotary_dim/2, rotary_dim/2), k_rot.narrow(-1, 0, rotary_dim/2)}, -1);
-        q_rot.mul_(cos).add_(rotate_half_q * sin);
-        k_rot.mul_(cos).add_(rotate_half_k * sin);
-        cos = at::Tensor(); sin = at::Tensor();
+// Selected dynamic pipeline requests must agree across every axis that can
+// execute the same adapter projection. PP stages intentionally have different
+// tensor layouts, so this uses the canonical request identity rather than the
+// stage-local adapter tensor layout hash.
+static void validate_pipeline_selected_adapter_request(
+    TrainingContext* ctx,
+    const int64_t* adapter_ids,
+    int32_t adapter_count
+) {
+    TORCH_CHECK(ctx && adapter_ids && adapter_count > 0,
+        "selected pipeline request requires at least one adapter ID");
+    AdapterRegistryHash hash;
+    hash_collective_request_topology(hash, ctx);
+    hash.add_u64(static_cast<uint64_t>(adapter_count));
+    int32_t found_count = 0;
+    for (int32_t index = 0; index < adapter_count; ++index) {
+        const int64_t requested_id = adapter_ids[index];
+        hash.add_u64(static_cast<uint64_t>(requested_id));
+        auto it = std::find_if(ctx->adapters.begin(), ctx->adapters.end(),
+            [requested_id](const auto& adapter) {
+                return adapter.id == requested_id;
+            });
+        if (it == ctx->adapters.end()) {
+            hash.add_u64(0x6d697373696e67ULL);
+            continue;
+        }
+        hash.add_u64(0x666f756e64ULL);
+        hash_adapter_request_identity(hash, *it);
+        ++found_count;
+    }
+    constexpr uint64_t kPositiveInt64Mask =
+        static_cast<uint64_t>(std::numeric_limits<int64_t>::max());
+    const std::vector<int64_t> values{
+        found_count == adapter_count ? 1 : 0,
+        adapter_count,
+        found_count,
+        static_cast<int64_t>(hash.first & kPositiveInt64Mask),
+        static_cast<int64_t>(hash.second & kPositiveInt64Mask),
+    };
+    auto options = at::TensorOptions().dtype(at::kLong).device(
+        at::kCUDA, ctx->cuda_device);
+    auto minimum = at::tensor(values, options);
+    auto maximum = minimum.clone();
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+    auto reduce_axis = [&](ncclComm_t communicator, const char* axis) {
+        if (!communicator) return;
+        auto error = ncclAllReduce(
+            minimum.data_ptr<int64_t>(), minimum.data_ptr<int64_t>(),
+            minimum.numel(), ncclInt64, ncclMin, communicator, stream);
+        TORCH_CHECK(error == ncclSuccess,
+            "selected pipeline request minimum ", axis,
+            " reduction failed: ", ncclGetErrorString(error));
+        error = ncclAllReduce(
+            maximum.data_ptr<int64_t>(), maximum.data_ptr<int64_t>(),
+            maximum.numel(), ncclInt64, ncclMax, communicator, stream);
+        TORCH_CHECK(error == ncclSuccess,
+            "selected pipeline request maximum ", axis,
+            " reduction failed: ", ncclGetErrorString(error));
+    };
+    reduce_axis(ctx->nccl_comm, "EP");
+    reduce_axis(ctx->dp_comm, "DP");
+    reduce_axis(ctx->tp_comm, "TP");
+    reduce_axis(ctx->cp_comm, "CP");
+    reduce_axis(ctx->pp_control_comm ? ctx->pp_control_comm : ctx->pp_comm,
+        "PP");
+    const auto minimum_cpu = minimum.to(at::kCPU);
+    const auto maximum_cpu = maximum.to(at::kCPU);
+    const auto* minimum_data = minimum_cpu.data_ptr<int64_t>();
+    const auto* maximum_data = maximum_cpu.data_ptr<int64_t>();
+    for (int64_t index = 0; index < minimum.numel(); ++index) {
+        TORCH_CHECK(minimum_data[index] == maximum_data[index],
+            "selected dynamic pipeline adapter request differs across "
+            "distributed ranks");
     }
+    TORCH_CHECK(minimum_data[0] == 1,
+        "selected dynamic pipeline adapter request is invalid on one or "
+        "more distributed ranks");
+}
 
-    // GQA: no K/V expansion needed (PT 2.5+ enable_gqa=true)
-    double scale = 1.0 / std::sqrt((double)head_dim);
+static bool adapter_collective_all_succeeded(
+    TrainingContext* ctx,
+    bool local_success
+) {
+    if (!ctx || (!ctx->nccl_comm && !ctx->dp_comm && !ctx->tp_comm &&
+            !ctx->cp_comm && !ctx->pp_control_comm && !ctx->pp_comm))
+        return local_success;
+    c10::cuda::set_device(ctx->cuda_device);
+    cudaSetDevice(ctx->cuda_device);
+    auto options = at::TensorOptions().dtype(at::kInt).device(
+        at::kCUDA, ctx->cuda_device);
+    auto succeeded = at::full({1}, local_success ? 1 : 0, options);
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+    auto reduce_axis = [&](ncclComm_t communicator, const char* axis) {
+        if (!communicator) return;
+        const auto err = ncclAllReduce(
+            succeeded.data_ptr<int32_t>(), succeeded.data_ptr<int32_t>(),
+            1, ncclInt32, ncclMin, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " adapter success consensus failed: ", ncclGetErrorString(err));
+    };
+    reduce_axis(ctx->nccl_comm, "EP");
+    reduce_axis(ctx->dp_comm, "DP");
+    reduce_axis(ctx->tp_comm, "TP");
+    reduce_axis(ctx->cp_comm, "CP");
+    reduce_axis(
+        ctx->pp_control_comm ? ctx->pp_control_comm : ctx->pp_comm, "PP");
+    return succeeded.to(at::kCPU).item<int32_t>() != 0;
+}
 
-    // SDPA with GQA
-    at::Tensor attn_out;
-    if (attention_mask.defined() && attention_mask.numel() > 0) {
-        auto kpm = attention_mask.to(at::kBool);
-        while (kpm.dim() > 2) kpm = kpm.squeeze(0);
-        if (kpm.size(0) == 1) {
-            kpm = kpm.unsqueeze(1).unsqueeze(1).expand({batch, 1, 1, seq});
-        } else {
-            kpm = kpm.unsqueeze(1).unsqueeze(1);
-        }
-        // Combined causal + padding mask
-        auto causal = at::triu(at::ones({seq, seq}, at::TensorOptions().dtype(at::kBool).device(q_out.device())), 1);
-        causal = causal.unsqueeze(0).unsqueeze(0);  // [1, 1, S, S]
-        auto pad_mask = kpm.logical_not();  // [B, 1, 1, S]
-        auto combined = causal.logical_or(pad_mask);
-        auto additive_mask = at::zeros({batch, 1, seq, seq}, at::TensorOptions().dtype(q_out.scalar_type()).device(q_out.device()));
-        additive_mask = additive_mask.masked_fill(combined, -std::numeric_limits<float>::infinity());
-        attn_out = at::scaled_dot_product_attention(q_out, k, v, additive_mask, 0.0, false, c10::nullopt, true);
-    } else {
-        attn_out = at::scaled_dot_product_attention(q_out, k, v, c10::nullopt, 0.0, true, c10::nullopt, true);
-    }
-    auto result = attn_out.transpose(1, 2).reshape({batch, seq, qkv_dim}).matmul(o_proj.t());
+static bool pipeline_control_all_succeeded(
+    TrainingContext* ctx,
+    bool local_success
+) {
+    if (!ctx || ctx->pp_world_size <= 1 ||
+            (!ctx->pp_control_comm && !ctx->pp_comm))
+        return local_success;
+    c10::cuda::set_device(ctx->cuda_device);
+    cudaSetDevice(ctx->cuda_device);
+    auto options = at::TensorOptions().dtype(at::kInt).device(
+        at::kCUDA, ctx->cuda_device);
+    auto status = at::full({1}, local_success ? 1 : 0, options);
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+    auto control_comm = ctx->pp_control_comm ?
+        ctx->pp_control_comm : ctx->pp_comm;
+    const auto error = ncclAllReduce(
+        status.data_ptr<int32_t>(), status.data_ptr<int32_t>(), 1,
+        ncclInt32, ncclMin, control_comm, stream);
+    TORCH_CHECK(error == ncclSuccess,
+        "PP control success consensus failed: ", ncclGetErrorString(error));
+    return status.to(at::kCPU).item<int32_t>() != 0;
+}
 
-    // Apply LoRA delta on o_proj output
-    auto it_o = ctx->lora_batch_cache.find(layer_idx * 10 + 3);
-    if (it_o != ctx->lora_batch_cache.end()) {
-        result = result + lora_activation_delta(attn_out.transpose(1, 2).reshape({batch, seq, qkv_dim}),
-            it_o->second.a_stack, it_o->second.b_stack, it_o->second.scaling);
+static bool fixed_optimizer_collective_all_succeeded(
+    TrainingContext* ctx,
+    bool local_success
+) {
+    bool succeeded = adapter_collective_all_succeeded(ctx, local_success);
+    return pipeline_control_all_succeeded(
+        ctx, succeeded && local_success);
+}
+
+static void validate_native_execution_topology_collective(
+    TrainingContext* ctx, bool allow_pipeline = false
+) {
+    std::string local_error;
+    try {
+        validate_native_execution_topology(ctx, allow_pipeline);
+    } catch (const std::exception& error) {
+        local_error = error.what();
     }
-    return result;
+    const bool local_valid = local_error.empty();
+    const bool globally_valid = adapter_collective_all_succeeded(
+        ctx, local_valid);
+    TORCH_CHECK(local_valid && globally_valid,
+        "native Qwen execution topology or model contract differs across "
+        "distributed ranks",
+        local_error.empty() ? "" : ": ", local_error);
 }
 
-static at::Tensor linear_attention_batched(
-    TrainingContext* ctx, const at::Tensor& hidden, int64_t layer_idx,
-    const at::Tensor& in_proj_qkv, const at::Tensor& in_proj_z,
-    const at::Tensor& in_proj_a, const at::Tensor& in_proj_b,
-    const at::Tensor& a_log, const at::Tensor& dt_bias,
-    const at::Tensor& conv1d_w, const at::Tensor& norm_w, const at::Tensor& out_proj,
-    int64_t num_k_heads, int64_t key_dim, int64_t num_v_heads, int64_t val_dim,
-    int64_t conv_kernel, double rms_eps, at::ScalarType compute_type
+static std::vector<int32_t> adapter_collective_min_flags(
+    TrainingContext* ctx,
+    const std::vector<int32_t>& local_flags
 ) {
-    // Full reimplementation of linear_attention non-chunked path with
-    // activation-level LoRA delta on QKV, Z, and out_proj.
-    auto device = hidden.device();
-    int64_t batch = hidden.size(0), seq = hidden.size(1);
-    int64_t q_size = num_k_heads * key_dim;
-    int64_t v_size = num_v_heads * val_dim;
-    int64_t qkv_dim = q_size * 2 + v_size;
+    TORCH_CHECK(!local_flags.empty(),
+        "adapter flag consensus requires at least one flag");
+    if (!ctx || (!ctx->nccl_comm && !ctx->dp_comm && !ctx->tp_comm &&
+            !ctx->cp_comm))
+        return local_flags;
+    c10::cuda::set_device(ctx->cuda_device);
+    cudaSetDevice(ctx->cuda_device);
+    auto flags = at::tensor(
+        local_flags, at::TensorOptions().dtype(at::kInt).device(
+            at::kCUDA, ctx->cuda_device));
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+    auto reduce_axis = [&](ncclComm_t communicator, const char* axis) {
+        if (!communicator) return;
+        const auto error = ncclAllReduce(
+            flags.data_ptr<int32_t>(), flags.data_ptr<int32_t>(),
+            flags.numel(), ncclInt32, ncclMin, communicator, stream);
+        TORCH_CHECK(error == ncclSuccess, axis,
+            " packed adapter flag consensus failed: ",
+            ncclGetErrorString(error));
+    };
+    reduce_axis(ctx->nccl_comm, "EP");
+    reduce_axis(ctx->dp_comm, "DP");
+    reduce_axis(ctx->tp_comm, "TP");
+    reduce_axis(ctx->cp_comm, "CP");
+    auto cpu = flags.to(at::kCPU).contiguous();
+    const auto* values = cpu.data_ptr<int32_t>();
+    return std::vector<int32_t>(values, values + cpu.numel());
+}
 
-    // QKV projection + LoRA delta
-    auto qkv = at::matmul(hidden, in_proj_qkv.t());
-    auto it_qkv = ctx->lora_batch_cache.find(layer_idx * 10 + 0);
-    if (it_qkv != ctx->lora_batch_cache.end()) {
-        qkv = qkv + lora_activation_delta(hidden, it_qkv->second.a_stack, it_qkv->second.b_stack, it_qkv->second.scaling);
-    }
+static void validate_dynamic_context_health(TrainingContext* ctx) {
+    TORCH_CHECK(ctx, "dynamic multi-LoRA requires a valid training context");
+    const bool local_healthy = !ctx->poisoned;
+    const bool globally_healthy = adapter_collective_all_succeeded(
+        ctx, local_healthy);
+    if (!globally_healthy) ctx->poisoned = true;
+    TORCH_CHECK(local_healthy && globally_healthy,
+        "dynamic multi-LoRA context is poisoned; restart the worker");
+}
 
-    // DIAG: dump after QKV projection
-    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
-        auto qkv_f = qkv.to(at::kFloat);
-        fprintf(stderr, "[diag-la] layer %ld qkv_proj: shape=[%ld,%ld,%ld] mean=%.6f std=%.6f [0,0,:5]=%.6f,%.6f,%.6f,%.6f,%.6f\n",
-                (long)layer_idx, (long)qkv_f.size(0), (long)qkv_f.size(1), (long)qkv_f.size(2),
-                qkv_f.mean().item<float>(), qkv_f.std().item<float>(),
-                qkv_f[0][0][0].item<float>(), qkv_f[0][0][1].item<float>(),
-                qkv_f[0][0][2].item<float>(), qkv_f[0][0][3].item<float>(),
-                qkv_f[0][0][4].item<float>());
-        // Dump weight layout info
-        auto w_f = in_proj_qkv.to(at::kFloat);
-        fprintf(stderr, "[diag-la] in_proj_qkv weight: shape=[%ld,%ld] mean=%.6f std=%.6f\n",
-                (long)w_f.size(0), (long)w_f.size(1), w_f.mean().item<float>(), w_f.std().item<float>());
-        auto conv_w_f = conv1d_w.to(at::kFloat);
-        fprintf(stderr, "[diag-la] conv1d weight: shape=[%ld,%ld,%ld] mean=%.6f std=%.6f\n",
-                (long)conv_w_f.size(0), (long)conv_w_f.size(1), (long)conv_w_f.size(2),
-                conv_w_f.mean().item<float>(), conv_w_f.std().item<float>());
-    }
+static void require_clear_accumulation_for_registry_mutation(
+    TrainingContext* ctx
+) {
+    validate_dynamic_context_health(ctx);
+    const bool local_clear = ctx && !ctx->accumulation_active &&
+        ctx->accumulated_token_weight == 0.0 &&
+        !ctx->pp_window.active;
+    const bool globally_clear = adapter_collective_all_succeeded(
+        ctx, local_clear);
+    TORCH_CHECK(local_clear && globally_clear,
+        "cannot mutate the dynamic LoRA registry while a gradient or "
+        "pipeline window is pending; finalize or abort it first");
+}
 
-    auto qkv_t = qkv.transpose(1, 2);
-    int64_t pad = conv_kernel - 1;
-    auto padding = at::zeros({batch, qkv_dim, pad}, qkv.options());
-    auto padded = at::cat({padding, qkv_t}, 2);
-    auto conv_out = at::conv1d(padded, conv1d_w, /*bias=*/{},
-        at::IntArrayRef({1}), at::IntArrayRef({0}), at::IntArrayRef({1}), qkv_dim);
-    conv_out = at::silu(conv_out.narrow(2, 0, seq));
-    auto qkv_conv = conv_out.transpose(1, 2);
+static bool adapter_registration_phase_matches(
+    TrainingContext* ctx,
+    const TrainingContext::LoRAAdapter& candidate,
+    bool local_ready,
+    uint64_t phase
+) {
+    if (!ctx || (!ctx->nccl_comm && !ctx->dp_comm && !ctx->tp_comm &&
+            !ctx->cp_comm && !ctx->pp_comm))
+        return local_ready;
+
+    AdapterRegistryHash hash;
+    hash_collective_topology(hash, ctx);
+    hash.add_u64(phase);
+    hash.add_u64(ctx->restore_without_parameter_sync);
+    hash.add_u64(ctx->allow_heterogeneous_registration);
+    hash.add_u64(ctx->next_adapter_id);
+    hash.add_u64(ctx->adapters.size());
+    for (const auto& adapter : ctx->adapters)
+        hash_adapter_layout(hash, adapter);
+    hash_adapter_layout(hash, candidate);
+
+    constexpr uint64_t kPositiveInt64Mask =
+        static_cast<uint64_t>(std::numeric_limits<int64_t>::max());
+    const std::vector<int64_t> signature_values{
+        local_ready ? 1 : 0,
+        static_cast<int64_t>(hash.first & kPositiveInt64Mask),
+        static_cast<int64_t>(hash.second & kPositiveInt64Mask),
+    };
+    c10::cuda::set_device(ctx->cuda_device);
+    cudaSetDevice(ctx->cuda_device);
+    auto options = at::TensorOptions().dtype(at::kLong).device(
+        at::kCUDA, ctx->cuda_device);
+    auto minimum = at::tensor(signature_values, options);
+    auto maximum = minimum.clone();
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+    auto reduce_axis = [&](ncclComm_t communicator, const char* axis) {
+        if (!communicator) return;
+        auto err = ncclAllReduce(
+            minimum.data_ptr<int64_t>(), minimum.data_ptr<int64_t>(),
+            minimum.numel(), ncclInt64, ncclMin, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " adapter registration minimum all-reduce failed: ",
+            ncclGetErrorString(err));
+        err = ncclAllReduce(
+            maximum.data_ptr<int64_t>(), maximum.data_ptr<int64_t>(),
+            maximum.numel(), ncclInt64, ncclMax, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " adapter registration maximum all-reduce failed: ",
+            ncclGetErrorString(err));
+    };
+    reduce_axis(ctx->nccl_comm, "EP");
+    reduce_axis(ctx->dp_comm, "DP");
+    reduce_axis(ctx->tp_comm, "TP");
+
+    const auto minimum_cpu = minimum.to(at::kCPU);
+    const auto maximum_cpu = maximum.to(at::kCPU);
+    const auto* minimum_data = minimum_cpu.data_ptr<int64_t>();
+    const auto* maximum_data = maximum_cpu.data_ptr<int64_t>();
+    bool stage_matches = minimum_data[0] == 1;
+    for (int64_t index = 0; index < minimum.numel(); ++index)
+        stage_matches = stage_matches &&
+            minimum_data[index] == maximum_data[index];
+
+    // Pipeline stages intentionally own different parameter layouts. Compare
+    // only the canonical tenant request across PP/CP while the full layout
+    // comparison above remains scoped to ranks sharing stage ownership.
+    AdapterRegistryHash request_hash;
+    hash_collective_request_topology(request_hash, ctx);
+    request_hash.add_u64(phase);
+    request_hash.add_u64(ctx->restore_without_parameter_sync);
+    request_hash.add_u64(ctx->allow_heterogeneous_registration);
+    request_hash.add_u64(ctx->next_adapter_id);
+    request_hash.add_u64(ctx->adapters.size());
+    for (const auto& adapter : ctx->adapters)
+        hash_adapter_request_identity(request_hash, adapter);
+    hash_adapter_request_identity(request_hash, candidate);
+    const std::vector<int64_t> request_values{
+        local_ready && stage_matches ? 1 : 0,
+        static_cast<int64_t>(request_hash.first & kPositiveInt64Mask),
+        static_cast<int64_t>(request_hash.second & kPositiveInt64Mask),
+    };
+    auto request_minimum = at::tensor(request_values, options);
+    auto request_maximum = request_minimum.clone();
+    auto reduce_request_axis = [&](ncclComm_t communicator, const char* axis) {
+        if (!communicator) return;
+        auto err = ncclAllReduce(
+            request_minimum.data_ptr<int64_t>(),
+            request_minimum.data_ptr<int64_t>(), request_minimum.numel(),
+            ncclInt64, ncclMin, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " adapter request minimum all-reduce failed: ",
+            ncclGetErrorString(err));
+        err = ncclAllReduce(
+            request_maximum.data_ptr<int64_t>(),
+            request_maximum.data_ptr<int64_t>(), request_maximum.numel(),
+            ncclInt64, ncclMax, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " adapter request maximum all-reduce failed: ",
+            ncclGetErrorString(err));
+    };
+    reduce_request_axis(ctx->cp_comm, "CP");
+    reduce_request_axis(ctx->pp_comm, "PP");
+    const auto request_minimum_cpu = request_minimum.to(at::kCPU);
+    const auto request_maximum_cpu = request_maximum.to(at::kCPU);
+    const auto* request_minimum_data =
+        request_minimum_cpu.data_ptr<int64_t>();
+    const auto* request_maximum_data =
+        request_maximum_cpu.data_ptr<int64_t>();
+    for (int64_t index = 0; index < request_minimum.numel(); ++index) {
+        if (request_minimum_data[index] != request_maximum_data[index])
+            return false;
+    }
+    return request_minimum_data[0] == 1;
+}
 
-    // DIAG: dump after conv1d
-    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
-        auto qc_f = qkv_conv.to(at::kFloat);
-        fprintf(stderr, "[diag-la] layer %ld after_conv1d: mean=%.6f std=%.6f [0,0,:5]=%.6f,%.6f,%.6f,%.6f,%.6f\n",
-                (long)layer_idx, qc_f.mean().item<float>(), qc_f.std().item<float>(),
-                qc_f[0][0][0].item<float>(), qc_f[0][0][1].item<float>(),
-                qc_f[0][0][2].item<float>(), qc_f[0][0][3].item<float>(),
-                qc_f[0][0][4].item<float>());
+static bool supported_mask_dtype(const at::Tensor& tensor) {
+    switch (tensor.scalar_type()) {
+        case at::kBool:
+        case at::kByte:
+        case at::kChar:
+        case at::kShort:
+        case at::kInt:
+        case at::kLong:
+        case at::kHalf:
+        case at::kBFloat16:
+        case at::kFloat:
+        case at::kDouble:
+            return true;
+        default:
+            return false;
     }
+}
 
-    // Flat QKV split (matches transformers Qwen3_5GatedDeltaNet.forward)
-    // in_proj_qkv outputs flat layout: [Q_all(2048) | K_all(2048) | V_all(4096)]
-    // NOT per-head interleaved. This matches the non-batched path (line ~517).
-    int64_t head_k_dim = key_dim;                        // 128 (already per-head)
-    int64_t head_v_dim = val_dim;                        // 128 (already per-head)
-    int64_t q_total = num_k_heads * head_k_dim;          // 2048
-    int64_t v_total = num_v_heads * head_v_dim;          // 4096
-    auto q = qkv_conv.narrow(-1, 0, q_total).reshape({batch, seq, num_k_heads, head_k_dim});
-    auto k = qkv_conv.narrow(-1, q_total, q_total).reshape({batch, seq, num_k_heads, head_k_dim});
-    auto v = qkv_conv.narrow(-1, q_total * 2, v_total).reshape({batch, seq, num_v_heads, head_v_dim});
+static bool replica_input_signatures_match(
+    TrainingContext* ctx,
+    const at::Tensor& input_ids,
+    const at::Tensor& target_mask,
+    const at::Tensor* attention_mask
+) {
+    if (!ctx || (!ctx->tp_comm && !ctx->nccl_comm && !ctx->cp_comm))
+        return true;
+    const std::vector<int64_t> signature_values{
+        input_ids.size(0),
+        input_ids.size(1),
+        static_cast<int64_t>(input_ids.scalar_type()),
+        target_mask.size(0),
+        target_mask.size(1),
+        static_cast<int64_t>(target_mask.scalar_type()),
+        attention_mask ? 1 : 0,
+        attention_mask ? attention_mask->size(0) : 0,
+        attention_mask ? attention_mask->size(1) : 0,
+        attention_mask
+            ? static_cast<int64_t>(attention_mask->scalar_type())
+            : 0,
+    };
+    auto options = at::TensorOptions().dtype(at::kLong).device(
+        at::kCUDA, ctx->cuda_device);
+    auto minimum = at::tensor(signature_values, options);
+    auto maximum = minimum.clone();
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+    auto reduce_axis = [&](ncclComm_t communicator, const char* axis) {
+        if (!communicator) return;
+        auto err = ncclAllReduce(
+            minimum.data_ptr<int64_t>(), minimum.data_ptr<int64_t>(),
+            minimum.numel(), ncclInt64, ncclMin, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " input signature minimum all-reduce failed: ",
+            ncclGetErrorString(err));
+        err = ncclAllReduce(
+            maximum.data_ptr<int64_t>(), maximum.data_ptr<int64_t>(),
+            maximum.numel(), ncclInt64, ncclMax, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, axis,
+            " input signature maximum all-reduce failed: ",
+            ncclGetErrorString(err));
+    };
+    reduce_axis(ctx->tp_comm, "TP");
+    reduce_axis(ctx->cp_comm, "CP");
+    const bool sharded_a2a = ctx->expert_parallel && ctx->nccl_comm &&
+        env_enabled("QWEN36_EP_A2A_SHARDED");
+    if (ctx->expert_parallel && !sharded_a2a)
+        reduce_axis(ctx->nccl_comm, "replicated EP");
+    const bool signatures_match =
+        (minimum == maximum).all().item<bool>();
+    if (!signatures_match || ctx->cp_world_size <= 1)
+        return signatures_match;
+
+    auto cp_values_match = [&](const at::Tensor& tensor, bool integer) {
+        auto values = tensor.to(integer ? at::kLong : at::kFloat)
+            .contiguous();
+        auto values_minimum = values.clone();
+        auto values_maximum = values.clone();
+        const auto dtype = integer ? ncclInt64 : ncclFloat;
+        auto error = ncclAllReduce(
+            values_minimum.data_ptr(), values_minimum.data_ptr(),
+            values_minimum.numel(), dtype, ncclMin, ctx->cp_comm, stream);
+        TORCH_CHECK(error == ncclSuccess,
+            "CP input-value minimum all-reduce failed: ",
+            ncclGetErrorString(error));
+        error = ncclAllReduce(
+            values_maximum.data_ptr(), values_maximum.data_ptr(),
+            values_maximum.numel(), dtype, ncclMax, ctx->cp_comm, stream);
+        TORCH_CHECK(error == ncclSuccess,
+            "CP input-value maximum all-reduce failed: ",
+            ncclGetErrorString(error));
+        return values_minimum.equal(values_maximum);
+    };
+    if (!cp_values_match(input_ids, true) ||
+            !cp_values_match(target_mask, false))
+        return false;
+    return !attention_mask || cp_values_match(*attention_mask, false);
+}
 
-    // DIAG: dump Q/K/V after per-head split
-    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
-        auto q_f = q.to(at::kFloat);
-        auto k_f = k.to(at::kFloat);
-        auto v_f = v.to(at::kFloat);
-        fprintf(stderr, "[diag-la] after_split q: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
-                q_f.mean().item<float>(), q_f.std().item<float>(),
-                q_f[0][0][0][0].item<float>(), q_f[0][0][0][1].item<float>(), q_f[0][0][0][2].item<float>());
-        fprintf(stderr, "[diag-la] after_split k: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
-                k_f.mean().item<float>(), k_f.std().item<float>(),
-                k_f[0][0][0][0].item<float>(), k_f[0][0][0][1].item<float>(), k_f[0][0][0][2].item<float>());
-        fprintf(stderr, "[diag-la] after_split v: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
-                v_f.mean().item<float>(), v_f.std().item<float>(),
-                v_f[0][0][0][0].item<float>(), v_f[0][0][0][1].item<float>(), v_f[0][0][0][2].item<float>());
+static bool harvest_leaf_grad(
+    at::Tensor& param, at::Tensor& fp32_accumulator
+) {
+    auto grad = param.grad();
+    if (!grad.defined()) return false;
+    TORCH_CHECK(fp32_accumulator.defined(),
+        "active LoRA parameter is missing its FP32 gradient accumulator");
+    TORCH_CHECK(fp32_accumulator.scalar_type() == at::kFloat,
+        "LoRA gradient accumulator must be FP32");
+    TORCH_CHECK(fp32_accumulator.sizes() == param.sizes(),
+        "LoRA gradient accumulator shape mismatch");
+    at::NoGradGuard guard;
+    fp32_accumulator.add_(grad.to(at::kFloat));
+    // Never retain BF16 leaf gradients across micro-batches. The next
+    // backward starts with a fresh leaf grad and is harvested again.
+    param.mutable_grad() = at::Tensor();
+    return true;
+}
+
+static bool harvest_adapter_gradients(TrainingContext::LoRAAdapter& adapter) {
+    bool harvested = false;
+    for (auto& [layer_idx, pairs] : adapter.params) {
+        auto accum_it = adapter.grad_accum.find(layer_idx);
+        TORCH_CHECK(accum_it != adapter.grad_accum.end() &&
+                    accum_it->second.size() == pairs.size(),
+            "dynamic LoRA gradient accumulator layout mismatch");
+        for (size_t i = 0; i < pairs.size(); ++i) {
+            auto& [a, b] = pairs[i];
+            auto& accum = accum_it->second[i];
+            if (a.requires_grad()) harvested |= harvest_leaf_grad(a, accum[0]);
+            if (b.requires_grad()) harvested |= harvest_leaf_grad(b, accum[1]);
+        }
     }
+    return harvested;
+}
 
-    auto a = at::matmul(hidden, in_proj_a.t());
-    auto b = at::matmul(hidden, in_proj_b.t());
+static bool harvest_gradient_accumulators(TrainingContext* ctx) {
+    bool harvested = false;
+    for (auto& adapter : ctx->adapters) {
+        harvested |= harvest_adapter_gradients(adapter);
+    }
+    TORCH_CHECK(ctx->grad_accum_a.size() == ctx->lora_a.size() &&
+                ctx->grad_accum_b.size() == ctx->lora_b.size(),
+        "fixed LoRA gradient accumulator layout mismatch");
+    for (size_t i = 0; i < ctx->lora_a.size(); ++i) {
+        if (!ctx->lora_active[i]) continue;
+        harvested |= harvest_leaf_grad(ctx->lora_a[i], ctx->grad_accum_a[i]);
+        harvested |= harvest_leaf_grad(ctx->lora_b[i], ctx->grad_accum_b[i]);
+    }
+    ctx->accumulation_active |= harvested;
+    return harvested;
+}
 
-    // Z projection + LoRA delta
-    auto z = at::matmul(hidden, in_proj_z.t());
-    auto it_z = ctx->lora_batch_cache.find(layer_idx * 10 + 1);
-    if (it_z != ctx->lora_batch_cache.end()) {
-        z = z + lora_activation_delta(hidden, it_z->second.a_stack, it_z->second.b_stack, it_z->second.scaling);
+static bool fixed_lora_accumulators_are_finite(TrainingContext* ctx) {
+    try {
+        TORCH_CHECK(ctx, "fixed LoRA finite check requires a training context");
+        if (ctx->fixed_grad_slab.defined())
+            return at::isfinite(ctx->fixed_grad_slab).all().item<bool>();
+        for (size_t i = 0; i < ctx->lora_a.size(); ++i) {
+            if (!ctx->lora_active[i]) continue;
+            for (const auto* accumulator : {
+                    &ctx->grad_accum_a[i], &ctx->grad_accum_b[i]}) {
+                if (!accumulator->defined() ||
+                        !at::isfinite(*accumulator).all().item<bool>())
+                    return false;
+            }
+        }
+        return true;
+    } catch (...) {
+        return false;
     }
-    z = z.reshape({batch, seq, num_v_heads, head_v_dim});
+}
 
-    // g = -exp(A_log) * softplus(a + dt_bias)
-    auto a_log_f = a_log.to(at::kFloat);
-    auto dt_bias_f = dt_bias.to(at::kFloat);
-    auto a_f = a.to(at::kFloat);
-    auto g = a_log_f.unsqueeze(0).unsqueeze(0).exp().neg() * at::softplus(a_f + dt_bias_f.unsqueeze(0).unsqueeze(0));
-    auto beta = at::sigmoid(b);
+static bool dynamic_lora_accumulators_are_finite(
+    TrainingContext* ctx,
+    const std::vector<uint8_t>& adapter_has_global_tokens
+) {
+    TORCH_CHECK(ctx && adapter_has_global_tokens.size() == ctx->adapters.size(),
+        "dynamic LoRA finite check requires matching adapter activity");
+    for (size_t adapter_index = 0;
+         adapter_index < ctx->adapters.size(); ++adapter_index) {
+        if (!adapter_has_global_tokens[adapter_index]) continue;
+        const auto& adapter = ctx->adapters[adapter_index];
+        if (adapter.grad_slab.defined()) {
+            if (!at::isfinite(adapter.grad_slab).all().item<bool>())
+                return false;
+            continue;
+        }
+        for (const auto& [layer_idx, pairs] : adapter.params) {
+            const auto accum_it = adapter.grad_accum.find(layer_idx);
+            if (accum_it == adapter.grad_accum.end() ||
+                    accum_it->second.size() != pairs.size())
+                return false;
+            for (size_t pair = 0; pair < pairs.size(); ++pair) {
+                if (!pairs[pair].first.requires_grad()) continue;
+                const auto& accumulators = accum_it->second[pair];
+                for (const auto& accumulator : accumulators) {
+                    if (!accumulator.defined() ||
+                            !at::isfinite(accumulator).all().item<bool>())
+                        return false;
+                }
+            }
+        }
+    }
+    return true;
+}
 
-    // Expand Q/K to num_v_heads
-    int64_t n_rep = num_v_heads / num_k_heads;
-    q = q.repeat_interleave(n_rep, 2);
-    k = k.repeat_interleave(n_rep, 2);
+static void clear_adapter_gradient_accumulators(
+    TrainingContext::LoRAAdapter& adapter
+) {
+    at::NoGradGuard guard;
+    if (adapter.grad_slab.defined()) adapter.grad_slab.zero_();
+    for (auto& [layer_idx, pairs] : adapter.params) {
+        auto accum_it = adapter.grad_accum.find(layer_idx);
+        for (size_t i = 0; i < pairs.size(); ++i) {
+            auto& [a, b] = pairs[i];
+            if (a.grad().defined()) a.mutable_grad() = at::Tensor();
+            if (b.grad().defined()) b.mutable_grad() = at::Tensor();
+            if (accum_it != adapter.grad_accum.end() &&
+                i < accum_it->second.size() && !adapter.grad_slab.defined()) {
+                if (accum_it->second[i][0].defined()) accum_it->second[i][0].zero_();
+                if (accum_it->second[i][1].defined()) accum_it->second[i][1].zero_();
+            }
+        }
+    }
+}
 
-    // L2 normalize Q, K (per-head, matching HF)
-    q = (q.to(at::kFloat) / q.to(at::kFloat).norm(2, -1, true).clamp_min(1e-6));
-    k = (k.to(at::kFloat) / k.to(at::kFloat).norm(2, -1, true).clamp_min(1e-6));
+static void clear_gradient_accumulators(TrainingContext* ctx) {
+    if (!ctx) return;
+    at::NoGradGuard guard;
+    for (auto& adapter : ctx->adapters) {
+        clear_adapter_gradient_accumulators(adapter);
+    }
+    if (ctx->fixed_grad_slab.defined()) ctx->fixed_grad_slab.zero_();
+    for (size_t i = 0; i < ctx->lora_a.size(); ++i) {
+        if (ctx->lora_a[i].grad().defined())
+            ctx->lora_a[i].mutable_grad() = at::Tensor();
+        if (ctx->lora_b[i].grad().defined())
+            ctx->lora_b[i].mutable_grad() = at::Tensor();
+        if (!ctx->fixed_grad_slab.defined() &&
+            i < ctx->grad_accum_a.size() && ctx->grad_accum_a[i].defined())
+            ctx->grad_accum_a[i].zero_();
+        if (!ctx->fixed_grad_slab.defined() &&
+            i < ctx->grad_accum_b.size() && ctx->grad_accum_b[i].defined())
+            ctx->grad_accum_b[i].zero_();
+    }
+    ctx->group_inputs.clear();
+    ctx->group_outputs.clear();
+    ctx->lora_cache_valid = false;
+    ctx->lora_batch_valid = false;
+    ctx->accumulation_active = false;
+    ctx->accumulated_token_weight = 0.0;
+}
 
-    // DIAG: dump after L2 norm
-    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
-        auto q_f = q.to(at::kFloat);
-        auto k_f = k.to(at::kFloat);
-        fprintf(stderr, "[diag-la] after_l2norm q: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
-                q_f.mean().item<float>(), q_f.std().item<float>(),
-                q_f[0][0][0][0].item<float>(), q_f[0][0][0][1].item<float>(), q_f[0][0][0][2].item<float>());
-        fprintf(stderr, "[diag-la] after_l2norm k: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
-                k_f.mean().item<float>(), k_f.std().item<float>(),
-                k_f[0][0][0][0].item<float>(), k_f[0][0][0][1].item<float>(), k_f[0][0][0][2].item<float>());
-        auto g_f = g.to(at::kFloat);
-        fprintf(stderr, "[diag-la] g: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
-                g_f.mean().item<float>(), g_f.std().item<float>(),
-                g_f[0][0][0].item<float>(), g_f[0][0][1].item<float>(), g_f[0][0][2].item<float>());
-        auto beta_f = beta.to(at::kFloat);
-        fprintf(stderr, "[diag-la] beta: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
-                beta_f.mean().item<float>(), beta_f.std().item<float>(),
-                beta_f[0][0][0].item<float>(), beta_f[0][0][1].item<float>(), beta_f[0][0][2].item<float>());
+struct GradientAccumulationFailureGuard {
+    TrainingContext* ctx;
+    bool disarmed = false;
+    ~GradientAccumulationFailureGuard() noexcept {
+        if (disarmed) return;
+        try {
+            clear_gradient_accumulators(ctx);
+        } catch (const std::exception& e) {
+            fprintf(stderr,
+                "[q36] secondary gradient cleanup failure: %s\n", e.what());
+        } catch (...) {
+            fprintf(stderr,
+                "[q36] secondary gradient cleanup failure: unknown error\n");
+        }
     }
+};
 
-    double scale = 1.0 / std::sqrt((double)head_k_dim);
-    q = q * scale;
+static at::Tensor tp_allreduce_lora_delta(
+    TrainingContext* ctx, const at::Tensor& local_delta
+) {
+    if (!ctx || ctx->tp_world_size <= 1) return local_delta;
+    TORCH_CHECK(ctx->tp_comm,
+        "LoRA TP communicator is not initialized for TP_SIZE=",
+        ctx->tp_world_size);
+    return NcclAllReduceFunction::apply(
+        local_delta, (int64_t)ctx->tp_comm,
+        (int64_t)reinterpret_cast<uintptr_t>(ctx->tp_stream));
+}
 
-    auto q_t = q.transpose(1, 2).contiguous();
-    auto k_t = k.transpose(1, 2).contiguous();
-    auto v_t = v.to(at::kFloat).transpose(1, 2).contiguous();
-    auto g_t = g.transpose(1, 2).contiguous();
-    auto beta_t = beta.to(at::kFloat).transpose(1, 2).contiguous();
+static at::Tensor tp_copy_lora_input(
+    TrainingContext* ctx, const at::Tensor& input
+) {
+    if (!ctx || ctx->tp_world_size <= 1) return input;
+    TORCH_CHECK(ctx->tp_comm,
+        "LoRA TP communicator is not initialized for TP_SIZE=",
+        ctx->tp_world_size);
+    return TpCopyToRegionFunction::apply(
+        input, (int64_t)ctx->tp_comm,
+        (int64_t)reinterpret_cast<uintptr_t>(ctx->tp_stream));
+}
 
-    auto g_exp = g_t.exp();
+static bool base_tp_attention_enabled(const TrainingContext* ctx) {
+    return ctx && ctx->base_tp_attention && ctx->tp_world_size > 1;
+}
 
-    // N-aware sub-batching: process N adapters in groups of 256 to limit
-    // state tensor size and improve SM occupancy.
-    // State is independent per adapter — safe to split by N dimension.
-    int64_t BH_total = batch * num_v_heads;
-    int64_t sub_batch = (BH_total > 8192) ? 256 : batch;  // 256 adapters or all if small
-    sub_batch = std::min(sub_batch, batch);
+static bool sequence_parallel_enabled(const TrainingContext* ctx) {
+    return ctx && ctx->sequence_parallel && ctx->tp_world_size > 1;
+}
 
-    auto outs = at::empty({BH_total, seq, head_v_dim}, q_t.options());
-    auto delta_buf = at::empty({BH_total, seq, head_v_dim}, q_t.options());
+static bool context_parallel_enabled(const TrainingContext* ctx) {
+    return ctx && ctx->cp_world_size > 1;
+}
 
-    for (int64_t sb = 0; sb < batch; sb += sub_batch) {
-        int64_t n = std::min(sub_batch, batch - sb);
-        int64_t BH = n * num_v_heads;
+static at::Tensor sequence_parallel_embedding_scatter(
+    TrainingContext* ctx, const at::Tensor& full_hidden
+) {
+    if (!sequence_parallel_enabled(ctx)) return full_hidden;
+    TORCH_CHECK(full_hidden.dim() == 3 &&
+            full_hidden.size(1) % ctx->tp_world_size == 0,
+        "sequence length=", full_hidden.dim() == 3 ? full_hidden.size(1) : -1,
+        " must be divisible by TP_SIZE=", ctx->tp_world_size);
+    const int64_t local_sequence = full_hidden.size(1) / ctx->tp_world_size;
+    ++ctx->sequence_parallel_embedding_scatter_count;
+    ctx->sequence_parallel_last_local_sequence = local_sequence;
+    return full_hidden.narrow(
+        1, ctx->tp_rank * local_sequence, local_sequence).contiguous();
+}
 
-        auto state = at::zeros({BH, head_k_dim, head_v_dim}, q_t.options());
+static at::Tensor sequence_parallel_column_gather(
+    TrainingContext* ctx, const at::Tensor& local_hidden
+) {
+    TORCH_CHECK(sequence_parallel_enabled(ctx) && ctx->tp_comm,
+        "sequence column gather requires an initialized TP communicator");
+    ++ctx->sequence_parallel_all_gather_count;
+    return TpGatherFromSequenceFunction::apply(
+        local_hidden.contiguous(), (int64_t)ctx->tp_comm,
+        (int64_t)reinterpret_cast<uintptr_t>(ctx->tp_stream),
+        ctx->tp_world_size);
+}
 
-        // Narrow on dim 0 (batch/adapter dimension), then reshape to [BH, seq, dim]
-        auto q_sub = q_t.narrow(0, sb, n);
-        auto k_sub = k_t.narrow(0, sb, n);
-        auto v_sub = v_t.narrow(0, sb, n);
-        auto g_sub = g_exp.narrow(0, sb, n);
-        auto beta_sub = beta_t.narrow(0, sb, n);
+// Full-attention CP keeps local queries and gathers the normalized, RoPE'd
+// K/V sequence once. The reduce-scatter backward is required because every
+// query rank contributes gradients to every global key/value owner.
+static at::Tensor context_parallel_kv_gather(
+    TrainingContext* ctx, const at::Tensor& local_kv
+) {
+    TORCH_CHECK(context_parallel_enabled(ctx) && ctx->cp_comm &&
+            env_enabled("QWEN36_CP_FULL_ATTENTION_KV_GATHER"),
+        "full-attention context parallelism requires the guarded KV gather");
+    return TpGatherFromSequenceFunction::apply(
+        local_kv.contiguous(), (int64_t)ctx->cp_comm,
+        (int64_t)reinterpret_cast<uintptr_t>(ctx->cp_stream),
+        ctx->cp_world_size);
+}
 
-        auto q_contig = q_sub.reshape({BH, seq, head_k_dim}).contiguous().to(at::kFloat);
-        auto k_contig = k_sub.reshape({BH, seq, head_k_dim}).contiguous().to(at::kFloat);
-        auto v_contig = v_sub.reshape({BH, seq, head_v_dim}).contiguous().to(at::kFloat);
-        auto g_contig = g_sub.reshape({BH, seq}).contiguous().to(at::kFloat);
-        auto beta_contig = beta_sub.reshape({BH, seq}).contiguous().to(at::kFloat);
-        auto state_contig = state.contiguous();
-
-        // Output slices: [BH, seq, val_dim] → need to narrow on dim 0 too
-        // outs is [BH_total, seq, val_dim], so bh_start = sb * num_v_heads
-        int64_t bh_start = sb * num_v_heads;
-        cuda_gated_delta_rule(
-            q_contig.data_ptr<float>(),
-            k_contig.data_ptr<float>(),
-            v_contig.data_ptr<float>(),
-            g_contig.data_ptr<float>(),
-            beta_contig.data_ptr<float>(),
-            state_contig.data_ptr<float>(),
-            outs.narrow(0, bh_start, BH).data_ptr<float>(),
-            delta_buf.narrow(0, bh_start, BH).data_ptr<float>(),
-            (int)BH, (int)seq, (int)head_k_dim, (int)head_v_dim
-        );
+static at::Tensor sequence_parallel_row_reduce_scatter(
+    TrainingContext* ctx, const at::Tensor& full_local_output
+) {
+    TORCH_CHECK(sequence_parallel_enabled(ctx) && ctx->tp_comm,
+        "sequence row reduce-scatter requires an initialized TP communicator");
+    ++ctx->sequence_parallel_reduce_scatter_count;
+    return TpReduceScatterToSequenceFunction::apply(
+        full_local_output, (int64_t)ctx->tp_comm,
+        (int64_t)reinterpret_cast<uintptr_t>(ctx->tp_stream),
+        ctx->tp_world_size);
+}
+
+static at::Tensor sequence_parallel_loss_gather(
+    TrainingContext* ctx, const at::Tensor& local_hidden
+) {
+    if (!sequence_parallel_enabled(ctx) && !context_parallel_enabled(ctx))
+        return local_hidden;
+    if (context_parallel_enabled(ctx)) {
+        return TpGatherForLossFunction::apply(
+            local_hidden.contiguous(), (int64_t)ctx->cp_comm,
+            (int64_t)reinterpret_cast<uintptr_t>(ctx->cp_stream),
+            ctx->cp_rank, ctx->cp_world_size);
     }
+    ++ctx->sequence_parallel_loss_gather_count;
+    return TpGatherForLossFunction::apply(
+        local_hidden.contiguous(), (int64_t)ctx->tp_comm,
+        (int64_t)reinterpret_cast<uintptr_t>(ctx->tp_stream),
+        ctx->tp_rank, ctx->tp_world_size);
+}
 
-    auto core_out = outs.reshape({batch, num_v_heads, seq, head_v_dim})
-                         .transpose(1, 2).to(compute_type);
+static at::Tensor sequence_parallel_plain_split(
+    TrainingContext* ctx, const at::Tensor& full_gradient
+) {
+    if (!sequence_parallel_enabled(ctx) && !context_parallel_enabled(ctx))
+        return full_gradient;
+    const int64_t world = context_parallel_enabled(ctx)
+        ? ctx->cp_world_size : ctx->tp_world_size;
+    const int64_t rank = context_parallel_enabled(ctx)
+        ? ctx->cp_rank : ctx->tp_rank;
+    TORCH_CHECK(full_gradient.dim() == 3 &&
+            full_gradient.size(1) % world == 0,
+        "loss hidden gradient sequence must be divisible by the sequence axis");
+    const int64_t local_sequence = full_gradient.size(1) / world;
+    return full_gradient.narrow(
+        1, rank * local_sequence, local_sequence).contiguous();
+}
 
-    // DIAG: dump after delta rule
-    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
-        auto co_f = core_out.to(at::kFloat);
-        fprintf(stderr, "[diag-la] after_delta_rule: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
-                co_f.mean().item<float>(), co_f.std().item<float>(),
-                co_f[0][0][0][0].item<float>(), co_f[0][0][0][1].item<float>(), co_f[0][0][0][2].item<float>());
-    }
+static bool vocab_parallel_enabled(const TrainingContext* ctx) {
+    return ctx && ctx->vocab_parallel && ctx->tp_world_size > 1;
+}
 
-    auto core_flat = core_out.reshape({-1, head_v_dim});
-    auto z_flat = z.reshape({-1, head_v_dim});
-    auto variance = core_flat.to(at::kFloat).pow(2).mean(-1, true);
-    auto normed = (core_flat.to(at::kFloat) * (variance + rms_eps).rsqrt() * norm_w.to(at::kFloat)).to(core_flat.scalar_type());
-    auto gated = (normed * at::silu(z_flat.to(at::kFloat)).to(normed.scalar_type())).reshape({batch, seq, num_v_heads * head_v_dim});
-    auto result = at::matmul(gated, out_proj.t());
+static at::Tensor tp_allreduce_value(
+    TrainingContext* ctx, const at::Tensor& input, ncclRedOp_t reduction
+) {
+    TORCH_CHECK(ctx && ctx->tp_comm,
+        "vocabulary TP communicator is not initialized for TP_SIZE=",
+        ctx ? ctx->tp_world_size : 0);
+    return NcclAllReduceFunction::allreduce(
+        input, ctx->tp_comm, ctx->tp_stream, reduction);
+}
 
-    // DIAG: dump after norm+gate+out_proj
-    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
-        auto r_f = result.to(at::kFloat);
-        fprintf(stderr, "[diag-la] after_out_proj: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
-                r_f.mean().item<float>(), r_f.std().item<float>(),
-                r_f[0][0][0].item<float>(), r_f[0][0][1].item<float>(), r_f[0][0][2].item<float>());
+static at::Tensor vocabulary_embedding(
+    TrainingContext* ctx, const at::Tensor& input_ids
+) {
+    auto embed = *ctx->embed_ptr[0];
+    if (!vocab_parallel_enabled(ctx)) {
+        if (context_parallel_enabled(ctx)) {
+            TORCH_CHECK(input_ids.dim() == 2 &&
+                    input_ids.size(1) % ctx->cp_world_size == 0,
+                "input sequence length=",
+                input_ids.dim() == 2 ? input_ids.size(1) : -1,
+                " must be divisible by CP_SIZE=", ctx->cp_world_size);
+            const int64_t local_sequence =
+                input_ids.size(1) / ctx->cp_world_size;
+            auto local_ids = input_ids.narrow(
+                1, ctx->cp_rank * local_sequence, local_sequence);
+            return at::embedding(embed, local_ids);
+        }
+        auto hidden = sequence_parallel_embedding_scatter(
+            ctx, at::embedding(embed, input_ids));
+        return hidden;
     }
 
-    // out_proj LoRA delta: result += B@(A@gated) * scaling
-    auto it_op = ctx->lora_batch_cache.find(layer_idx * 10 + 2);
-    if (it_op != ctx->lora_batch_cache.end()) {
-        result = result + lora_activation_delta(gated, it_op->second.a_stack, it_op->second.b_stack, it_op->second.scaling);
-    }
+    TORCH_CHECK(!context_parallel_enabled(ctx),
+        "vocabulary TP cannot be combined with GDN context parallelism");
+
+    const int64_t vocab_start = ctx->tp_rank * ctx->local_vocab_size;
+    const int64_t vocab_end = vocab_start + ctx->local_vocab_size;
+    auto in_range = (input_ids >= vocab_start) & (input_ids < vocab_end);
+    auto local_ids = (input_ids - vocab_start).clamp(0, ctx->local_vocab_size - 1);
+    auto local_hidden = at::embedding(embed, local_ids);
+    local_hidden = local_hidden * in_range.unsqueeze(-1).to(local_hidden.scalar_type());
+    return sequence_parallel_embedding_scatter(
+        ctx, tp_allreduce_value(ctx, local_hidden, ncclSum));
+}
 
-    return result;
+static bool base_tp_mlp_enabled(const TrainingContext* ctx) {
+    return ctx && ctx->base_tp_mlp && ctx->tp_world_size > 1;
 }
-// ──────────────────────────────────────────────────────────────────────
 
-static at::Tensor dense_mlp_forward(
-    const at::Tensor& hidden,
-    const at::Tensor& gate_proj, const at::Tensor& up_proj, const at::Tensor& down_proj,
-    at::ScalarType compute_type
+static int64_t base_tp_mlp_world_size(const TrainingContext* ctx) {
+    TORCH_CHECK(base_tp_mlp_enabled(ctx),
+        "base MLP TP context is not enabled");
+    return ctx->tp_world_size;
+}
+
+static at::Tensor tp_allreduce_base_mlp(
+    TrainingContext* ctx, const at::Tensor& local_output
 ) {
-    int64_t batch = hidden.size(0), seq = hidden.size(1), hidden_dim = hidden.size(2);
-    auto flat = hidden.reshape({batch * seq, hidden_dim});
-    auto gate_out = at::matmul(flat, gate_proj.t());
-    auto up_out = at::matmul(flat, up_proj.t());
-    // Fused silu * up via Tilelang (falls back to ATen)
-    auto activated = fused_swiglu_op(gate_out, up_out, 0.0);  // Qwen3.6 dense MLP has no clamp
-    return at::matmul(activated, down_proj.t()).reshape({batch, seq, hidden_dim});
+    if (!ctx || !ctx->base_tp_mlp || ctx->tp_world_size <= 1)
+        return local_output;
+    TORCH_CHECK(ctx->tp_comm,
+        "base MLP TP communicator is not initialized for TP_SIZE=",
+        ctx->tp_world_size);
+    if (sequence_parallel_enabled(ctx))
+        return sequence_parallel_row_reduce_scatter(ctx, local_output);
+    return NcclAllReduceFunction::apply(
+        local_output, (int64_t)ctx->tp_comm,
+        (int64_t)reinterpret_cast<uintptr_t>(ctx->tp_stream));
 }
 
-// Forward pass (no checkpointing)
-static at::Tensor forward_full(
-    TrainingContext* ctx,
-    const at::Tensor& input_ids
+static at::Tensor tp_copy_base_mlp_input(
+    TrainingContext* ctx, const at::Tensor& input
 ) {
-    if (ctx->lora_batch_valid) prepare_lora_batch(ctx);
-    else precompute_lora_cache(ctx);
-    auto kind = ctx->compute_type;
-    auto embed = *ctx->embed_ptr[0];
-    auto final_norm = *ctx->final_norm_ptr[0];
+    if (!ctx || !ctx->base_tp_mlp || ctx->tp_world_size <= 1)
+        return input;
+    TORCH_CHECK(ctx->tp_comm,
+        "base MLP TP communicator is not initialized for TP_SIZE=",
+        ctx->tp_world_size);
+    if (sequence_parallel_enabled(ctx))
+        return sequence_parallel_column_gather(ctx, input);
+    return TpCopyToRegionFunction::apply(
+        input, (int64_t)ctx->tp_comm,
+        (int64_t)reinterpret_cast<uintptr_t>(ctx->tp_stream));
+}
 
-    at::AutoGradMode guard(true);
-    at::Tensor hidden = at::embedding(embed, input_ids);
+static at::Tensor tp_allreduce_base_attention(
+    TrainingContext* ctx, const at::Tensor& local_output
+) {
+    if (!ctx || !ctx->base_tp_attention || ctx->tp_world_size <= 1)
+        return local_output;
+    TORCH_CHECK(ctx->tp_comm,
+        "base attention TP communicator is not initialized for TP_SIZE=",
+        ctx->tp_world_size);
+    if (sequence_parallel_enabled(ctx))
+        return sequence_parallel_row_reduce_scatter(ctx, local_output);
+    return NcclAllReduceFunction::apply(
+        local_output, (int64_t)ctx->tp_comm,
+        (int64_t)reinterpret_cast<uintptr_t>(ctx->tp_stream));
+}
 
-    // Debug: dump embedding output stats
-    if (getenv("QWEN36_DUMP_LAYERS")) {
-        auto h_f = hidden.to(at::kFloat);
-        fprintf(stderr, "[dump] embedding: mean=%.6f std=%.6f [0,:3]=%.6f,%.6f,%.6f\n",
-                h_f.mean().item<float>(), h_f.std().item<float>(),
-                h_f[0][0][0].item<float>(), h_f[0][0][1].item<float>(), h_f[0][0][2].item<float>());
-    }
+static at::Tensor tp_copy_base_attention_input(
+    TrainingContext* ctx, const at::Tensor& input
+) {
+    if (!ctx || !ctx->base_tp_attention || ctx->tp_world_size <= 1)
+        return input;
+    TORCH_CHECK(ctx->tp_comm,
+        "base attention TP communicator is not initialized for TP_SIZE=",
+        ctx->tp_world_size);
+    if (sequence_parallel_enabled(ctx))
+        return sequence_parallel_column_gather(ctx, input);
+    return TpCopyToRegionFunction::apply(
+        input, (int64_t)ctx->tp_comm,
+        (int64_t)reinterpret_cast<uintptr_t>(ctx->tp_stream));
+}
 
-    for (int64_t i = 0; i < ctx->num_layers; i++) {
-        // Get weight pointers for this layer
-        int64_t w_offset = 0;
-        for (int64_t j = 0; j < i; j++)
-            w_offset += weight_count_for_layer(ctx->layer_configs[j]);
-        int64_t w_count = weight_count_for_layer(ctx->layer_configs[i]);
-        std::vector<at::Tensor*> layer_w(ctx->weight_ptrs.begin() + w_offset,
-                                         ctx->weight_ptrs.begin() + w_offset + w_count);
+static LoraTpLayout lora_tp_layout(
+    const TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx
+) {
+    if (!ctx || ctx->tp_world_size <= 1 ||
+        layer_idx < 0 || layer_idx >= ctx->num_layers)
+        return LoraTpLayout::LatentRank;
+    const auto& cfg = ctx->layer_configs[layer_idx];
+    auto table = lora_projection_table(cfg);
+    TORCH_CHECK(pair_idx >= 0 && pair_idx < table.count,
+        "invalid LoRA pair for TP layout");
+    const std::string name(table.entries[pair_idx].name);
+    if (ctx->base_tp_attention &&
+        (name == "q_proj" || name == "k_proj" || name == "v_proj" ||
+         name == "in_proj_qkv" || name == "in_proj_z" ||
+         name == "in_proj_a" || name == "in_proj_b"))
+        return LoraTpLayout::ColumnParallel;
+    if (ctx->base_tp_attention &&
+        (name == "o_proj" || name == "out_proj"))
+        return LoraTpLayout::RowParallel;
+    if (ctx->base_tp_mlp &&
+        (name == "gate_proj" || name == "up_proj" ||
+         name == "shared_gate_proj" || name == "shared_up_proj" ||
+         name == "experts_gate_up_proj"))
+        return LoraTpLayout::ColumnParallel;
+    if (ctx->base_tp_mlp &&
+        (name == "down_proj" || name == "shared_down_proj" ||
+         name == "experts_down_proj"))
+        return LoraTpLayout::RowParallel;
+    return LoraTpLayout::LatentRank;
+}
 
-        // Get LoRA pointers for this layer (nullptr if no LoRA for this layer)
-        int64_t lora_count = (ctx->layer_configs[i].layer_type == 0) ? 4 : 3;
-        int64_t la_offset = ctx->lora_layer_offset[i];
-        // Check if this layer has LoRA params (la_offset < lora_a.size())
-        bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size();
-        std::vector<at::Tensor*> la_ptrs(lora_count, nullptr), lb_ptrs(lora_count, nullptr);
-        if (has_lora) {
-            for (int64_t k = 0; k < lora_count; k++) {
-                la_ptrs[k] = &ctx->lora_a[la_offset + k];
-                lb_ptrs[k] = &ctx->lora_b[la_offset + k];
-            }
-        }
+static bool lora_tp_parameter_replicated(LoraTpLayout layout, bool is_a) {
+    return (layout == LoraTpLayout::ColumnParallel && is_a) ||
+        (layout == LoraTpLayout::RowParallel && !is_a);
+}
 
-        hidden = forward_single_layer(ctx, hidden, layer_w.data(), &ctx->layer_configs[i], i,
-            kind, ctx->attention_mask, ctx->lora_batch_valid);
+struct LoraGradientSlabBinding {
+    at::Tensor* accumulator = nullptr;
+    const at::Tensor* parameter = nullptr;
+    uint8_t bucket_key = 0;
+    int64_t offset = 0;
+};
 
-        // Debug: dump per-layer hidden state stats
-        if (getenv("QWEN36_DUMP_LAYERS")) {
-            auto h_f = hidden.to(at::kFloat);
-            fprintf(stderr, "[dump] layer %ld: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
-                    (long)i, h_f.mean().item<float>(), h_f.std().item<float>(),
-                    h_f[0][0][0].item<float>(), h_f[0][0][1].item<float>(), h_f[0][0][2].item<float>());
+static uint8_t lora_gradient_bucket_key(
+    const TrainingContext* ctx, int64_t layer, int64_t pair, bool is_a
+) {
+    const auto table = lora_projection_table(ctx->layer_configs[layer]);
+    const bool grouped_expert = table.entries[pair].grouped_expert;
+    const auto layout = lora_tp_layout(ctx, layer, pair);
+    const bool tp_replicated =
+        (layout == LoraTpLayout::ColumnParallel && is_a) ||
+        (layout == LoraTpLayout::RowParallel && !is_a);
+    return (grouped_expert ? 2 : 0) | (tp_replicated ? 1 : 0);
+}
+
+static void bind_lora_gradient_slab(
+    std::vector<LoraGradientSlabBinding>& bindings,
+    at::Tensor& slab
+) {
+    if (bindings.empty()) {
+        slab = at::Tensor();
+        return;
+    }
+    constexpr int64_t alignment = 64;
+    int64_t total = 0;
+    for (uint8_t key = 0; key < 4; ++key) {
+        total = ((total + alignment - 1) / alignment) * alignment;
+        for (auto& binding : bindings) {
+            if (binding.bucket_key != key) continue;
+            binding.offset = total;
+            total += binding.parameter->numel();
         }
+    }
+    slab = at::zeros({total}, at::TensorOptions()
+        .dtype(at::kFloat).device(bindings.front().parameter->device()));
+    for (auto& binding : bindings) {
+        *binding.accumulator = slab.narrow(
+            0, binding.offset, binding.parameter->numel())
+            .view(binding.parameter->sizes());
+    }
+}
 
-        // No per-layer sync — let CUDA pipeline run asynchronously.
-        // emptyCache() here was the #1 cause of GPU underutilization (6% util).
+static void bind_fixed_lora_gradient_slab(TrainingContext* ctx) {
+    if (!env_enabled("QWEN36_GRAD_SLAB", true)) {
+        for (size_t index = 0; index < ctx->lora_a.size(); ++index) {
+            if (!ctx->lora_active[index]) continue;
+            const auto options = at::TensorOptions().dtype(at::kFloat)
+                .device(ctx->lora_a[index].device());
+            ctx->grad_accum_a[index] = at::zeros(
+                ctx->lora_a[index].sizes(), options);
+            ctx->grad_accum_b[index] = at::zeros(
+                ctx->lora_b[index].sizes(), options);
+        }
+        return;
+    }
+    std::vector<LoraGradientSlabBinding> bindings;
+    for (int64_t layer = 0; layer < ctx->num_layers; ++layer) {
+        const int64_t offset = ctx->lora_layer_offset[layer];
+        const int64_t count = lora_pair_count(ctx->layer_configs[layer]);
+        for (int64_t pair = 0; pair < count; ++pair) {
+            const int64_t index = offset + pair;
+            if (!ctx->lora_active[index]) continue;
+            bindings.push_back({
+                &ctx->grad_accum_a[index], &ctx->lora_a[index],
+                lora_gradient_bucket_key(ctx, layer, pair, true), 0});
+            bindings.push_back({
+                &ctx->grad_accum_b[index], &ctx->lora_b[index],
+                lora_gradient_bucket_key(ctx, layer, pair, false), 0});
+        }
     }
+    bind_lora_gradient_slab(bindings, ctx->fixed_grad_slab);
+}
 
-    return hidden;  // pre-norm hidden (for MTP)
+static void bind_adapter_lora_gradient_slab(
+    TrainingContext* ctx, TrainingContext::LoRAAdapter& adapter
+) {
+    if (!env_enabled("QWEN36_GRAD_SLAB", true)) {
+        for (auto& [layer, pairs] : adapter.params) {
+            auto& accumulators = adapter.grad_accum.at(layer);
+            for (size_t pair = 0; pair < pairs.size(); ++pair) {
+                auto& [a, b] = pairs[pair];
+                if (!a.requires_grad()) continue;
+                const auto options = at::TensorOptions().dtype(at::kFloat)
+                    .device(a.device());
+                accumulators[pair][0] = at::zeros(a.sizes(), options);
+                accumulators[pair][1] = at::zeros(b.sizes(), options);
+            }
+        }
+        return;
+    }
+    std::vector<LoraGradientSlabBinding> bindings;
+    for (auto& [layer, pairs] : adapter.params) {
+        auto accum_it = adapter.grad_accum.find(layer);
+        TORCH_CHECK(accum_it != adapter.grad_accum.end() &&
+                accum_it->second.size() == pairs.size(),
+            "dynamic LoRA gradient slab layout mismatch");
+        for (int64_t pair = 0; pair < static_cast<int64_t>(pairs.size()); ++pair) {
+            auto& [a, b] = pairs[pair];
+            if (!a.requires_grad()) continue;
+            bindings.push_back({
+                &accum_it->second[pair][0], &a,
+                lora_gradient_bucket_key(ctx, layer, pair, true), 0});
+            bindings.push_back({
+                &accum_it->second[pair][1], &b,
+                lora_gradient_bucket_key(ctx, layer, pair, false), 0});
+        }
+    }
+    bind_lora_gradient_slab(bindings, adapter.grad_slab);
 }
 
-// ──────────────────────────────────────────────────────────────────────
-// Gradient checkpointing: per-group recomputation
-// ──────────────────────────────────────────────────────────────────────
+// Dynamic tenants may be registered long before they are selected for a
+// training step. Keep their Adam state and transactional shadow handles lazy
+// in that case; the activation parameters and FP32 gradient slab remain ready
+// for the first selected backward. The helper is idempotent so checkpoint
+// hydration and a partially failed allocation can safely retry it.
+static void materialize_dynamic_adam_state(
+    TrainingContext::LoRAAdapter& adapter
+) {
+    at::NoGradGuard guard;
+    const bool host_state_current =
+        adapter.adam_host_step == adapter.optimizer_step;
+    for (auto& [layer_idx, pairs] : adapter.params) {
+        auto& states = adapter.adam_state[layer_idx];
+        auto& shadows = adapter.adam_shadow[layer_idx];
+        auto& shadow_slots = adapter.adam_shadow_slots[layer_idx];
+        TORCH_CHECK(states.size() == pairs.size() &&
+                shadows.size() == pairs.size() &&
+                shadow_slots.size() == pairs.size(),
+            "dynamic optimizer lazy-state layout mismatch for adapter ",
+            adapter.id, " layer ", layer_idx);
+        for (size_t pair_idx = 0; pair_idx < pairs.size(); ++pair_idx) {
+            auto& [a, b] = pairs[pair_idx];
+            if (!a.requires_grad()) continue;
+            auto& [m_a, v_a, m_b, v_b] = states[pair_idx];
+            auto f32_options = at::TensorOptions().dtype(at::kFloat)
+                .device(a.device());
+            const bool device_state_complete =
+                m_a.defined() && v_a.defined() &&
+                m_b.defined() && v_b.defined();
+            if (device_state_complete) continue;
+
+            const std::array<at::Tensor, 4>* host_state = nullptr;
+            if (host_state_current) {
+                auto host_it = adapter.adam_host_state.find(layer_idx);
+                TORCH_CHECK(host_it != adapter.adam_host_state.end() &&
+                        host_it->second.size() == pairs.size(),
+                    "dynamic Adam host-state layout mismatch for adapter ",
+                    adapter.id, " layer ", layer_idx);
+                host_state = &host_it->second[pair_idx];
+                for (const auto& tensor : *host_state) {
+                    TORCH_CHECK(tensor.defined() && tensor.device().is_cpu() &&
+                            tensor.is_contiguous() &&
+                            tensor.scalar_type() == at::kFloat,
+                        "dynamic Adam host state is incomplete for adapter ",
+                        adapter.id, " layer ", layer_idx,
+                        " pair ", pair_idx);
+                }
+                TORCH_CHECK((*host_state)[0].sizes() == a.sizes() &&
+                        (*host_state)[1].sizes() == a.sizes() &&
+                        (*host_state)[2].sizes() == b.sizes() &&
+                        (*host_state)[3].sizes() == b.sizes(),
+                    "dynamic Adam host-state shape mismatch for adapter ",
+                    adapter.id, " layer ", layer_idx,
+                    " pair ", pair_idx);
+            } else {
+                TORCH_CHECK(adapter.optimizer_step == 0,
+                    "trained dynamic adapter has neither device nor current "
+                    "host Adam state: adapter ", adapter.id,
+                    " step ", adapter.optimizer_step);
+            }
+            if (!m_a.defined()) m_a = host_state
+                ? (*host_state)[0].to(a.device())
+                : at::zeros(a.sizes(), f32_options);
+            if (!v_a.defined()) v_a = host_state
+                ? (*host_state)[1].to(a.device())
+                : at::zeros(a.sizes(), f32_options);
+            if (!m_b.defined()) m_b = host_state
+                ? (*host_state)[2].to(b.device())
+                : at::zeros(b.sizes(), f32_options);
+            if (!v_b.defined()) v_b = host_state
+                ? (*host_state)[3].to(b.device())
+                : at::zeros(b.sizes(), f32_options);
+            TORCH_CHECK(m_a.device() == a.device() &&
+                    v_a.device() == a.device() &&
+                    m_b.device() == b.device() &&
+                    v_b.device() == b.device() &&
+                    m_a.scalar_type() == at::kFloat &&
+                    v_a.scalar_type() == at::kFloat &&
+                    m_b.scalar_type() == at::kFloat &&
+                    v_b.scalar_type() == at::kFloat &&
+                    m_a.is_contiguous() && v_a.is_contiguous() &&
+                    m_b.is_contiguous() && v_b.is_contiguous() &&
+                    m_a.sizes() == a.sizes() &&
+                    v_a.sizes() == a.sizes() &&
+                    m_b.sizes() == b.sizes() &&
+                    v_b.sizes() == b.sizes(),
+                "dynamic Adam materialized-state mismatch for adapter ",
+                adapter.id, " layer ", layer_idx, " pair ", pair_idx);
+        }
+    }
+}
 
-// Run a group of layers forward (with grad enabled, for recomputation)
-static at::Tensor forward_layer_group(
+static void snapshot_dynamic_adam_state_to_host(
     TrainingContext* ctx,
-    const at::Tensor& input,
-    int64_t start_layer,
-    int64_t end_layer
+    TrainingContext::LoRAAdapter& adapter
 ) {
-    auto kind = ctx->compute_type;
-    at::Tensor hidden = input;
-
-    // Normal path: full layer forward (sub-layer checkpointing is handled
-    // in forward_full_checkpoint, not here)
-    for (int64_t i = start_layer; i < end_layer; i++) {
-        int64_t w_offset = 0;
-        for (int64_t j = 0; j < i; j++)
-            w_offset += weight_count_for_layer(ctx->layer_configs[j]);
-        int64_t w_count = weight_count_for_layer(ctx->layer_configs[i]);
-        std::vector<at::Tensor*> layer_w(ctx->weight_ptrs.begin() + w_offset,
-                                         ctx->weight_ptrs.begin() + w_offset + w_count);
-
-        int64_t lora_count = (ctx->layer_configs[i].layer_type == 0) ? 4 : 3;
-        int64_t la_offset = ctx->lora_layer_offset[i];
-        bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size();
-        std::vector<at::Tensor*> la_ptrs(lora_count, nullptr), lb_ptrs(lora_count, nullptr);
-        if (has_lora) {
-            for (int64_t k = 0; k < lora_count; k++) {
-                la_ptrs[k] = &ctx->lora_a[la_offset + k];
-                lb_ptrs[k] = &ctx->lora_b[la_offset + k];
+    TORCH_CHECK(ctx, "dynamic Adam host snapshot requires a context");
+    TORCH_CHECK(!ctx->dynamic_adam_transaction_active,
+        "cannot snapshot dynamic Adam during an active transaction");
+    TORCH_CHECK(adapter.optimizer_step > 0,
+        "untrained dynamic adapter has no Adam checkpoint state");
+    if (adapter.adam_host_step == adapter.optimizer_step) return;
+
+    const int64_t snapshot_step = adapter.optimizer_step;
+    // Production Adam commits are asynchronous. Checkpointing is an explicit
+    // persistence boundary and may run on a different host thread/current
+    // stream, so complete all device writes before starting D2H copies.
+    const auto sync_error = cudaDeviceSynchronize();
+    TORCH_CHECK(sync_error == cudaSuccess,
+        "dynamic Adam checkpoint device synchronization failed: ",
+        cudaGetErrorString(sync_error));
+    TORCH_CHECK(!ctx->dynamic_adam_transaction_active &&
+            adapter.optimizer_step == snapshot_step,
+        "dynamic Adam changed before its host snapshot copy");
+    std::map<int64_t, std::vector<std::array<at::Tensor, 4>>> snapshot;
+    at::NoGradGuard guard;
+    for (const auto& [layer_idx, pairs] : adapter.params) {
+        auto state_it = adapter.adam_state.find(layer_idx);
+        TORCH_CHECK(state_it != adapter.adam_state.end() &&
+                state_it->second.size() == pairs.size(),
+            "dynamic Adam checkpoint layout mismatch for adapter ",
+            adapter.id, " layer ", layer_idx);
+        std::vector<std::array<at::Tensor, 4>> host_states;
+        host_states.reserve(pairs.size());
+        for (size_t pair_idx = 0; pair_idx < pairs.size(); ++pair_idx) {
+            const auto& [a, b] = pairs[pair_idx];
+            if (!a.requires_grad()) {
+                host_states.push_back(std::array<at::Tensor, 4>{});
+                continue;
+            }
+            const auto& state = state_it->second[pair_idx];
+            const auto& m_a = state[0];
+            const auto& v_a = state[1];
+            const auto& m_b = state[2];
+            const auto& v_b = state[3];
+            for (const auto* tensor : {&m_a, &v_a, &m_b, &v_b}) {
+                TORCH_CHECK(tensor->defined() && tensor->is_cuda() &&
+                        tensor->is_contiguous() &&
+                        tensor->scalar_type() == at::kFloat,
+                    "dynamic Adam checkpoint requires contiguous CUDA FP32 "
+                    "state for adapter ", adapter.id, " layer ", layer_idx,
+                    " pair ", pair_idx);
             }
+            TORCH_CHECK(m_a.sizes() == a.sizes() &&
+                    v_a.sizes() == a.sizes() &&
+                    m_b.sizes() == b.sizes() &&
+                    v_b.sizes() == b.sizes() &&
+                    m_a.device() == a.device() &&
+                    v_a.device() == a.device() &&
+                    m_b.device() == b.device() &&
+                    v_b.device() == b.device(),
+                "dynamic Adam checkpoint shape mismatch for adapter ",
+                adapter.id, " layer ", layer_idx, " pair ", pair_idx);
+            host_states.push_back({
+                m_a.to(at::kCPU).contiguous(),
+                v_a.to(at::kCPU).contiguous(),
+                m_b.to(at::kCPU).contiguous(),
+                v_b.to(at::kCPU).contiguous()});
         }
+        snapshot.emplace(layer_idx, std::move(host_states));
+    }
+    TORCH_CHECK(!ctx->dynamic_adam_transaction_active &&
+            adapter.optimizer_step == snapshot_step,
+        "dynamic Adam changed while creating its host snapshot");
+    adapter.adam_host_state.swap(snapshot);
+    adapter.adam_host_step = snapshot_step;
+}
 
-        hidden = forward_single_layer(ctx, hidden, layer_w.data(), &ctx->layer_configs[i], i,
-            kind, ctx->attention_mask, ctx->lora_batch_valid);
-
-        // Debug: dump per-layer hidden state stats (also in checkpoint recompute path)
-        if (getenv("QWEN36_DUMP_LAYERS")) {
-            auto h_f = hidden.to(at::kFloat);
-            fprintf(stderr, "[dump] layer %ld: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
-                    (long)i, h_f.mean().item<float>(), h_f.std().item<float>(),
-                    h_f[0][0][0].item<float>(), h_f[0][0][1].item<float>(), h_f[0][0][2].item<float>());
+static uint64_t dynamic_adam_resident_bytes(
+    const TrainingContext::LoRAAdapter& adapter
+) {
+    uint64_t bytes = 0;
+    for (const auto& [layer_idx, states] : adapter.adam_state) {
+        (void)layer_idx;
+        for (const auto& state : states) {
+            for (const auto& tensor : state) {
+                if (!tensor.defined()) continue;
+                const uint64_t count = static_cast<uint64_t>(tensor.numel());
+                const uint64_t width =
+                    static_cast<uint64_t>(tensor.element_size());
+                if (width != 0 && count >
+                        (std::numeric_limits<uint64_t>::max() - bytes) / width)
+                    return std::numeric_limits<uint64_t>::max();
+                bytes += count * width;
+            }
         }
     }
-    return hidden;
+    return bytes;
 }
 
-// autograd::Function for checkpointing a group of layers.
-// Forward: run group WITHOUT grad (no intermediate activations stored).
-// Backward: recompute group WITH grad, then backprop through recomputed graph.
-struct GroupCheckpointFunction : public torch::autograd::Function<GroupCheckpointFunction> {
-    static at::Tensor forward(
-        torch::autograd::AutogradContext* ctx,
-        at::Tensor input,
-        int64_t tc_val,
-        int64_t start_layer,
-        int64_t end_layer
-    ) {
-        ctx->saved_data["tc"] = tc_val;
-        ctx->saved_data["start"] = start_layer;
-        ctx->saved_data["end"] = end_layer;
+static void clear_dynamic_adam_device_state(
+    TrainingContext::LoRAAdapter& adapter
+) {
+    std::map<int64_t, std::vector<std::array<at::Tensor, 4>>> empty_state;
+    for (const auto& [layer_idx, pairs] : adapter.params) {
+        empty_state.emplace(layer_idx,
+            std::vector<std::array<at::Tensor, 4>>(pairs.size()));
+    }
+    adapter.adam_state.swap(empty_state);
+}
 
-        // Check if activation offload is enabled
-        bool offload = getenv("QWEN36_OFFLOAD_ACTIVATIONS");
+static bool try_dynamic_adam_resident_budget(uint64_t* budget) {
+    if (!budget) return false;
+    const char* value = std::getenv(
+        "QWEN36_DYNAMIC_ADAM_RESIDENT_BYTES");
+    if (!value || value[0] == '\0' || value[0] == '-') return false;
+    errno = 0;
+    char* end = nullptr;
+    const unsigned long long parsed = std::strtoull(value, &end, 10);
+    if (errno == ERANGE || !end || *end != '\0') return false;
+    *budget = static_cast<uint64_t>(parsed);
+    return true;
+}
 
-        if (offload) {
-            // Save input to CPU — frees GPU memory between groups
-            // We store the CPU copy in saved_data (not save_for_backward,
-            // because save_for_backward would keep it on GPU)
-            auto input_cpu = input.detach().to(at::TensorOptions().dtype(input.scalar_type()).device(at::kCPU).pinned_memory(true));
-            ctx->saved_data["input_cpu"] = input_cpu;
-            // Also store device for restoring later
-            ctx->saved_data["device"] = input.device();
-        } else {
-            ctx->save_for_backward({input});
+// Paging is intentionally a post-commit, best-effort maintenance action. It
+// must never turn an already committed optimizer step into a reported failure.
+static void page_cold_dynamic_adam_state(
+    TrainingContext* ctx, const int64_t* protected_ids,
+    int32_t protected_count
+) noexcept {
+    if (!ctx || !env_enabled("QWEN36_DYNAMIC_ADAM_HOST_PAGING")) return;
+    uint64_t budget = 0;
+    if (!try_dynamic_adam_resident_budget(&budget)) {
+        fprintf(stderr,
+            "[q36] dynamic Adam paging ignored invalid or missing "
+            "QWEN36_DYNAMIC_ADAM_RESIDENT_BYTES\n");
+        return;
+    }
+    try {
+        TORCH_CHECK(!ctx->dynamic_adam_transaction_active,
+            "dynamic Adam paging requires a completed transaction");
+        if (ctx->dynamic_adam_lru_tick ==
+                std::numeric_limits<uint64_t>::max()) {
+            ctx->dynamic_adam_lru_tick = 0;
+            for (auto& adapter : ctx->adapters)
+                adapter.adam_last_used_tick = 0;
+        }
+        const uint64_t access_tick = ++ctx->dynamic_adam_lru_tick;
+        std::set<int64_t> protected_set;
+        for (int32_t index = 0; index < protected_count; ++index) {
+            if (protected_ids && protected_ids[index] > 0)
+                protected_set.insert(protected_ids[index]);
+        }
+        for (auto& adapter : ctx->adapters) {
+            if (protected_set.count(adapter.id) != 0)
+                adapter.adam_last_used_tick = access_tick;
         }
 
-        // Run forward in NO-GRAD mode — intermediate activations are NOT stored.
-        at::AutoGradMode guard(false);
-        auto* tc = reinterpret_cast<TrainingContext*>(tc_val);
-        return forward_layer_group(tc, input, start_layer, end_layer);
+        uint64_t resident = 0;
+        std::vector<TrainingContext::LoRAAdapter*> candidates;
+        for (auto& adapter : ctx->adapters) {
+            const uint64_t adapter_bytes =
+                dynamic_adam_resident_bytes(adapter);
+            if (adapter_bytes > std::numeric_limits<uint64_t>::max() - resident)
+                resident = std::numeric_limits<uint64_t>::max();
+            else
+                resident += adapter_bytes;
+            if (adapter_bytes > 0 &&
+                    protected_set.count(adapter.id) == 0)
+                candidates.push_back(&adapter);
+        }
+        if (resident <= budget) return;
+        std::sort(candidates.begin(), candidates.end(),
+            [](const auto* lhs, const auto* rhs) {
+                if (lhs->adam_last_used_tick != rhs->adam_last_used_tick)
+                    return lhs->adam_last_used_tick < rhs->adam_last_used_tick;
+                return lhs->id < rhs->id;
+            });
+        for (auto* adapter : candidates) {
+            if (resident <= budget) break;
+            const uint64_t adapter_bytes =
+                dynamic_adam_resident_bytes(*adapter);
+            if (adapter_bytes == 0) continue;
+            try {
+                if (adapter->optimizer_step > 0)
+                    snapshot_dynamic_adam_state_to_host(ctx, *adapter);
+                clear_dynamic_adam_device_state(*adapter);
+                resident = adapter_bytes >= resident
+                    ? 0
+                    : resident - adapter_bytes;
+            } catch (const std::exception& error) {
+                fprintf(stderr,
+                    "[q36] dynamic Adam paging kept adapter %ld resident: %s\n",
+                    static_cast<long>(adapter->id), error.what());
+            }
+        }
+    } catch (const std::exception& error) {
+        fprintf(stderr, "[q36] dynamic Adam paging skipped: %s\n",
+            error.what());
+    } catch (...) {
+        fprintf(stderr, "[q36] dynamic Adam paging skipped: unknown error\n");
     }
+}
 
-    static std::vector<at::Tensor> backward(
-        torch::autograd::AutogradContext* ctx,
-        std::vector<at::Tensor> grad_output
-    ) {
-        bool offload = ctx->saved_data.count("input_cpu") > 0;
+static bool dynamic_adam_shadow_layout_matches(
+    const std::array<at::Tensor, 6>& tensors,
+    const at::Tensor& a,
+    const at::Tensor& b
+) {
+    for (const auto& tensor : tensors) {
+        if (!tensor.defined() || !tensor.is_cuda() ||
+                !tensor.is_contiguous())
+            return false;
+    }
+    return tensors[0].sizes() == a.sizes() &&
+        tensors[0].device() == a.device() &&
+        tensors[0].scalar_type() == a.scalar_type() &&
+        tensors[0].requires_grad() == a.requires_grad() &&
+        tensors[1].sizes() == a.sizes() &&
+        tensors[1].device() == a.device() &&
+        tensors[1].scalar_type() == at::kFloat &&
+        tensors[2].sizes() == a.sizes() &&
+        tensors[2].device() == a.device() &&
+        tensors[2].scalar_type() == at::kFloat &&
+        tensors[3].sizes() == b.sizes() &&
+        tensors[3].device() == b.device() &&
+        tensors[3].scalar_type() == b.scalar_type() &&
+        tensors[3].requires_grad() == b.requires_grad() &&
+        tensors[4].sizes() == b.sizes() &&
+        tensors[4].device() == b.device() &&
+        tensors[4].scalar_type() == at::kFloat &&
+        tensors[5].sizes() == b.sizes() &&
+        tensors[5].device() == b.device() &&
+        tensors[5].scalar_type() == at::kFloat;
+}
 
-        at::Tensor input;
-        if (offload) {
-            // Restore input from CPU → GPU
-            auto input_cpu = ctx->saved_data["input_cpu"].toTensor();
-            auto device = ctx->saved_data["device"].toDevice();
-            input = input_cpu.to(device);
-        } else {
-            auto saved = ctx->get_saved_variables();
-            input = saved[0];
+static int64_t acquire_dynamic_adam_shadow_slot(
+    TrainingContext* ctx,
+    const at::Tensor& a,
+    const at::Tensor& b
+) {
+    TORCH_CHECK(ctx, "dynamic Adam shadow pool requires a context");
+    size_t reusable_index = ctx->dynamic_adam_shadow_pool.size();
+    for (size_t index = 0;
+         index < ctx->dynamic_adam_shadow_pool.size(); ++index) {
+        auto& entry = ctx->dynamic_adam_shadow_pool[index];
+        if (entry.in_use) continue;
+        if (dynamic_adam_shadow_layout_matches(entry.tensors, a, b)) {
+            entry.in_use = true;
+            return static_cast<int64_t>(index);
         }
+        if (reusable_index == ctx->dynamic_adam_shadow_pool.size())
+            reusable_index = index;
+    }
+    TORCH_CHECK(ctx->dynamic_adam_shadow_pool.size() <
+            static_cast<size_t>(std::numeric_limits<int64_t>::max()),
+        "dynamic Adam shadow pool exhausted its index range");
+    auto f32_options = at::TensorOptions().dtype(at::kFloat)
+        .device(a.device());
+    std::array<at::Tensor, 6> replacement{
+        at::empty_like(a).set_requires_grad(a.requires_grad()),
+        at::empty(a.sizes(), f32_options),
+        at::empty(a.sizes(), f32_options),
+        at::empty_like(b).set_requires_grad(b.requires_grad()),
+        at::empty(b.sizes(), f32_options),
+        at::empty(b.sizes(), f32_options)};
+    if (reusable_index == ctx->dynamic_adam_shadow_pool.size()) {
+        ctx->dynamic_adam_shadow_pool.emplace_back();
+        reusable_index = ctx->dynamic_adam_shadow_pool.size() - 1;
+    }
+    auto& entry = ctx->dynamic_adam_shadow_pool[reusable_index];
+    entry.tensors = std::move(replacement);
+    entry.in_use = true;
+    return static_cast<int64_t>(reusable_index);
+}
 
-        auto tc = reinterpret_cast<TrainingContext*>(ctx->saved_data["tc"].toInt());
-        int64_t start_layer = ctx->saved_data["start"].toInt();
-        int64_t end_layer = ctx->saved_data["end"].toInt();
-
-        // Recompute forward WITH grad enabled — builds autograd graph for LoRA params.
-        at::AutoGradMode guard(true);
-        input.set_requires_grad(true);
-        auto output = forward_layer_group(tc, input, start_layer, end_layer);
+static void materialize_dynamic_adam_shadow(
+    TrainingContext* ctx,
+    TrainingContext::LoRAAdapter& adapter
+) {
+    at::NoGradGuard guard;
+    const bool use_pool = env_enabled(
+        "QWEN36_DYNAMIC_ADAM_SHADOW_POOL", true);
+    for (auto& [layer_idx, pairs] : adapter.params) {
+        auto& states = adapter.adam_state[layer_idx];
+        auto& shadows = adapter.adam_shadow[layer_idx];
+        auto& shadow_slots = adapter.adam_shadow_slots[layer_idx];
+        TORCH_CHECK(states.size() == pairs.size() &&
+                shadows.size() == pairs.size() &&
+                shadow_slots.size() == pairs.size(),
+            "dynamic optimizer lazy-state layout mismatch for adapter ",
+            adapter.id, " layer ", layer_idx);
+        for (size_t pair_idx = 0; pair_idx < pairs.size(); ++pair_idx) {
+            auto& [a, b] = pairs[pair_idx];
+            if (!a.requires_grad()) continue;
+            auto& [next_a, next_m_a, next_v_a,
+                   next_b, next_m_b, next_v_b] = shadows[pair_idx];
+            if (use_pool) {
+                auto& slot = shadow_slots[pair_idx];
+                if (slot < 0) {
+                    slot = acquire_dynamic_adam_shadow_slot(ctx, a, b);
+                    shadows[pair_idx] =
+                        ctx->dynamic_adam_shadow_pool[slot].tensors;
+                }
+                TORCH_CHECK(slot < static_cast<int64_t>(
+                                ctx->dynamic_adam_shadow_pool.size()) &&
+                        ctx->dynamic_adam_shadow_pool[slot].in_use &&
+                        dynamic_adam_shadow_layout_matches(
+                            shadows[pair_idx], a, b),
+                    "dynamic Adam shadow pool lease is invalid for adapter ",
+                    adapter.id, " layer ", layer_idx,
+                    " pair ", pair_idx);
+                continue;
+            }
+            auto f32_options = at::TensorOptions().dtype(at::kFloat)
+                .device(a.device());
+            if (!next_a.defined())
+                next_a = at::empty_like(a).set_requires_grad(true);
+            if (!next_m_a.defined())
+                next_m_a = at::empty(a.sizes(), f32_options);
+            if (!next_v_a.defined())
+                next_v_a = at::empty(a.sizes(), f32_options);
+            if (!next_b.defined())
+                next_b = at::empty_like(b).set_requires_grad(true);
+            if (!next_m_b.defined())
+                next_m_b = at::empty(b.sizes(), f32_options);
+            if (!next_v_b.defined())
+                next_v_b = at::empty(b.sizes(), f32_options);
+        }
+    }
+}
 
-        // Backprop through recomputed graph.
-        // retain_graph=false: each group's recomputed graph is independent.
-        // LoRA param gradients accumulate via autograd's accumulator (leaf nodes).
-        // The graph is freed immediately after backward — critical for memory.
-        torch::autograd::backward({output}, {grad_output[0]},
-            /*retain_graph=*/false, /*create_graph=*/false);
-        return {input.grad(), at::Tensor(), at::Tensor(), at::Tensor()};
+static void release_dynamic_adam_shadows(
+    TrainingContext* ctx,
+    TrainingContext::LoRAAdapter& adapter
+) noexcept {
+    if (!ctx) return;
+    for (auto& [layer_idx, shadows] : adapter.adam_shadow) {
+        auto slots_it = adapter.adam_shadow_slots.find(layer_idx);
+        if (slots_it == adapter.adam_shadow_slots.end()) continue;
+        auto& slots = slots_it->second;
+        const size_t count = std::min(shadows.size(), slots.size());
+        for (size_t pair_idx = 0; pair_idx < count; ++pair_idx) {
+            const int64_t slot = slots[pair_idx];
+            if (slot < 0 || slot >= static_cast<int64_t>(
+                    ctx->dynamic_adam_shadow_pool.size()))
+                continue;
+            auto& entry = ctx->dynamic_adam_shadow_pool[slot];
+            entry.tensors = std::move(shadows[pair_idx]);
+            shadows[pair_idx] = std::array<at::Tensor, 6>{};
+            entry.in_use = false;
+            slots[pair_idx] = -1;
+        }
     }
-};
+}
 
-// ──────────────────────────────────────────────────────────────────────
-// FusedLayerFunction: autograd::Function for single layer forward+backward.
-// Forward: run WITH grad (PyTorch saves intermediates in graph).
-// Backward: PyTorch autograd traverses graph — NO recompute needed.
-// This eliminates checkpoint recompute (the main bottleneck).
-// Controlled by QWEN36_FUSED_LAYER=1 env var.
-// ──────────────────────────────────────────────────────────────────────
+static void release_dynamic_adam_shadows(
+    TrainingContext* ctx,
+    std::vector<TrainingContext::LoRAAdapter>& adapters
+) noexcept {
+    if (!ctx) return;
+    for (auto& adapter : adapters)
+        release_dynamic_adam_shadows(ctx, adapter);
+}
 
-struct FusedLayerFunction : public torch::autograd::Function<FusedLayerFunction> {
-    static at::Tensor forward(
-        torch::autograd::AutogradContext* ctx,
-        at::Tensor input,
-        int64_t tc_val,
-        int64_t layer_idx
-    ) {
-        ctx->saved_data["tc"] = tc_val;
-        ctx->saved_data["layer"] = layer_idx;
-        // No save_for_backward — PyTorch autograd graph handles it.
-        // Forward runs WITH grad — all intermediates saved in graph.
-        auto* tc = reinterpret_cast<TrainingContext*>(tc_val);
-        auto kind = tc->compute_type;
+static void release_dynamic_adam_shadows(TrainingContext* ctx) noexcept {
+    if (!ctx) return;
+    release_dynamic_adam_shadows(ctx, ctx->adapters);
+}
 
-        int64_t w_offset = 0;
-        for (int64_t j = 0; j < layer_idx; j++)
-            w_offset += weight_count_for_layer(tc->layer_configs[j]);
-        int64_t w_count = weight_count_for_layer(tc->layer_configs[layer_idx]);
-        std::vector<at::Tensor*> layer_w(tc->weight_ptrs.begin() + w_offset,
-                                         tc->weight_ptrs.begin() + w_offset + w_count);
-        int64_t lora_count = (tc->layer_configs[layer_idx].layer_type == 0) ? 4 : 3;
-        int64_t la_offset = tc->lora_layer_offset[layer_idx];
-        bool has_lora = (la_offset + lora_count) <= (int64_t)tc->lora_a.size();
-        std::vector<at::Tensor*> la(lora_count, nullptr), lb(lora_count, nullptr);
-        if (has_lora) for (int64_t k = 0; k < lora_count; k++) {
-            la[k] = &tc->lora_a[la_offset + k]; lb[k] = &tc->lora_b[la_offset + k];
-        }
-        return forward_single_layer(tc, input, layer_w.data(), &tc->layer_configs[layer_idx],
-            layer_idx, kind, tc->attention_mask, tc->lora_batch_valid);
-    }
+struct DynamicAdamShadowLeaseScope {
+    TrainingContext* ctx = nullptr;
+    bool retain_for_outer_transaction = false;
+    bool active = false;
 
-    static std::vector<at::Tensor> backward(
-        torch::autograd::AutogradContext* ctx,
-        std::vector<at::Tensor> grad_output
-    ) {
-        // PyTorch autograd handles backward through the graph built during forward.
-        // No recompute needed — just return grad_output as grad_input.
-        // The actual backward computation happens via PyTorch's autograd engine
-        // traversing the graph nodes (matmul backward, SDPA backward, etc.).
-        return {grad_output[0], at::Tensor(), at::Tensor()};
+    ~DynamicAdamShadowLeaseScope() {
+        if (active && !retain_for_outer_transaction) {
+            release_dynamic_adam_shadows(ctx);
+            if (ctx) ctx->dynamic_adam_transaction_active = false;
+        }
     }
 };
 
-// Forward pass with fused layer (no checkpoint, no recompute).
-// Uses FusedLayerFunction per layer — PyTorch autograd handles backward.
-// QWEN36_FUSED_LAYER=1 enables this path.
-static at::Tensor forward_full_fused(
+static void materialize_dynamic_optimizer_state(
     TrainingContext* ctx,
-    const at::Tensor& input_ids
+    TrainingContext::LoRAAdapter& adapter
 ) {
-    if (ctx->lora_batch_valid) prepare_lora_batch(ctx);
-    else precompute_lora_cache(ctx);
-    auto embed = *ctx->embed_ptr[0];
-    at::Tensor hidden = at::embedding(embed, input_ids);
-    hidden = hidden.detach().set_requires_grad(true);
+    materialize_dynamic_adam_state(adapter);
+    materialize_dynamic_adam_shadow(ctx, adapter);
+}
 
-    for (int64_t i = 0; i < ctx->num_layers; i++) {
-        hidden = FusedLayerFunction::apply(
-            hidden,
-            (int64_t)(uintptr_t)ctx,
-            i
-        );
+static bool active_lora_targets_use_latent_rank_layout(
+    const TrainingContext* ctx,
+    const std::set<int64_t>& target_layers,
+    const std::set<std::string>& target_modules,
+    bool all_target_layers,
+    bool empty_modules_mean_attention_only
+) {
+    for (int64_t layer = 0; layer < ctx->num_layers; ++layer) {
+        if (!all_target_layers && target_layers.count(layer) == 0) continue;
+        const auto table = lora_projection_table(ctx->layer_configs[layer]);
+        for (int64_t pair = 0; pair < table.count; ++pair) {
+            const auto& projection = table.entries[pair];
+            const bool active = target_modules.empty()
+                ? !empty_modules_mean_attention_only ||
+                    (!projection.grouped_expert &&
+                     projection.segment == LoraSegment::Attention)
+                : target_modules.count(projection.name) > 0;
+            if (active &&
+                lora_tp_layout(ctx, layer, pair) == LoraTpLayout::LatentRank)
+                return true;
+        }
     }
+    return false;
+}
 
-    return hidden;
+static int64_t local_lora_rank_for_active_targets(
+    const TrainingContext* ctx,
+    int64_t global_rank,
+    const std::set<int64_t>& target_layers,
+    const std::set<std::string>& target_modules,
+    bool all_target_layers,
+    bool empty_modules_mean_attention_only,
+    const char* adapter_kind
+) {
+    TORCH_CHECK(global_rank > 0, adapter_kind, " LoRA rank must be positive");
+    const bool uses_latent_rank = active_lora_targets_use_latent_rank_layout(
+        ctx, target_layers, target_modules, all_target_layers,
+        empty_modules_mean_attention_only);
+    TORCH_CHECK(!uses_latent_rank || global_rank % ctx->tp_world_size == 0,
+        adapter_kind, " LoRA rank ", global_rank,
+        " must be divisible by TP_SIZE=", ctx->tp_world_size,
+        " because at least one active projection uses latent-rank sharding");
+    return uses_latent_rank ? global_rank / ctx->tp_world_size : global_rank;
 }
 
-// Forward pass with gradient checkpointing — manual checkpoint (no autograd::Function)
-// Forward: no-grad, save group inputs (offloaded to CPU). Backward: manual recompute per group.
-// This avoids autograd engine retaining all group outputs simultaneously.
-static at::Tensor forward_full_checkpoint(
-    TrainingContext* ctx,
-    const at::Tensor& input_ids
+static at::Tensor initialize_lora_a(
+    TrainingContext* ctx, const at::TensorOptions& options,
+    int64_t experts, int64_t global_rank, int64_t in_features
 ) {
-    // Use batched path if multiple adapters, else legacy weight-level
-    if (ctx->lora_batch_valid) {
-        prepare_lora_batch(ctx);
-    } else {
-        precompute_lora_cache(ctx);
+    const int64_t local_rank = global_rank / ctx->tp_world_size;
+    const int64_t rank_start = ctx->tp_rank * local_rank;
+    if (experts > 0) {
+        auto global = at::randn({experts, global_rank, in_features}, options);
+        return global.narrow(1, rank_start, local_rank).contiguous() * 0.01;
     }
-    auto embed = *ctx->embed_ptr[0];
-    at::Tensor hidden = at::embedding(embed, input_ids);
+    auto global = at::randn({global_rank, in_features}, options);
+    return global.narrow(0, rank_start, local_rank).contiguous() * 0.01;
+}
 
-    if (getenv("QWEN36_DUMP_LAYERS")) {
-        auto h_f = hidden.to(at::kFloat);
-        fprintf(stderr, "[dump] ckpt embedding: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
-                h_f.mean().item<float>(), h_f.std().item<float>(),
-                h_f[0][0][0].item<float>(), h_f[0][0][1].item<float>(), h_f[0][0][2].item<float>());
-    }
+static const char* lora_pair_name(const LayerConfig& cfg, int64_t pair_idx) {
+    auto table = lora_projection_table(cfg);
+    TORCH_CHECK(pair_idx >= 0 && pair_idx < table.count, "invalid LoRA projection index");
+    return table.entries[pair_idx].name;
+}
 
-    bool use_subckpt = getenv("QWEN36_SUBCKPT");
+static constexpr int64_t LORA_CACHE_STRIDE = 32;
 
-    if (use_subckpt) {
-        at::AutoGradMode restore(true);
-        hidden = hidden.detach().set_requires_grad(true);
-        for (int64_t i = 0; i < ctx->num_layers; i++) {
-            hidden = forward_single_layer_subckpt(ctx, hidden, i);
-        }
-        return hidden;
-    }
+static inline int64_t lora_cache_key(int64_t layer_idx, int64_t pair_idx) {
+    return layer_idx * LORA_CACHE_STRIDE + pair_idx;
+}
 
-    // Group-level manual checkpointing with variable group size.
-    // Larger group_size = fewer recomputations in backward (faster) but more
-    // peak memory. Default gs=4 (from ctx->group_size), overridable via env.
-    at::AutoGradMode no_grad(false);
-    hidden = hidden.detach();
+static inline bool legacy_lora_slot_active(const TrainingContext* ctx, int64_t slot) {
+    return slot >= 0 && slot < (int64_t)ctx->lora_active.size() && ctx->lora_active[slot] != 0;
+}
 
-    ctx->group_inputs.clear();
-    bool offload = getenv("QWEN36_OFFLOAD_ACTIVATIONS");
+static RoutedExpertLora routed_expert_lora(
+    TrainingContext* ctx, int64_t layer_idx, const LayerConfig& cfg
+) {
+    RoutedExpertLora result;
+    // Dynamic multi-LoRA supplies per-sample activation-level tensors. Do not
+    // mix the fixed adapter's expert tensors into those batches.
+    if (!ctx->adapters.empty() || cfg.num_experts <= 0 || layer_idx < 0 ||
+            layer_idx >= (int64_t)ctx->lora_layer_offset.size())
+        return result;
 
-    // Build group list using ctx->group_size (default 4).
-    // Env override: QWEN36_GROUP_SIZE=10 sets gs=10.
-    int64_t gs = ctx->group_size;
-    if (gs < 1) gs = 1;
-    const char* gs_env = getenv("QWEN36_GROUP_SIZE");
-    if (gs_env) { gs = atol(gs_env); if (gs < 1) gs = 1; }
-    fprintf(stderr, "[checkpoint] group_size=%ld (num_layers=%ld → %ld groups)\n",
-            (long)gs, (long)ctx->num_layers,
-            (long)((ctx->num_layers + gs - 1) / gs));
+    const int64_t offset = ctx->lora_layer_offset[layer_idx];
+    const int64_t gate_up_pair = lora_pair_index(cfg, "experts_gate_up_proj");
+    const int64_t down_pair = lora_pair_index(cfg, "experts_down_proj");
+    if (gate_up_pair >= 0 && legacy_lora_slot_active(ctx, offset + gate_up_pair)) {
+        result.gate_up_a = &ctx->lora_a[offset + gate_up_pair];
+        result.gate_up_b = &ctx->lora_b[offset + gate_up_pair];
+    }
+    if (down_pair >= 0 && legacy_lora_slot_active(ctx, offset + down_pair)) {
+        result.down_a = &ctx->lora_a[offset + down_pair];
+        result.down_b = &ctx->lora_b[offset + down_pair];
+    }
+    result.scaling = ctx->lora_scaling;
+    return result;
+}
 
-    std::vector<std::pair<int64_t, int64_t>> groups;
-    for (int64_t i = 0; i < ctx->num_layers; i += gs) {
-        groups.push_back({i, std::min(i + gs, ctx->num_layers)});
+static ncclDataType_t nccl_dtype_for(const at::Tensor& tensor) {
+    switch (tensor.scalar_type()) {
+        case at::kBFloat16: return ncclBfloat16;
+        case at::kFloat: return ncclFloat;
+        case at::kHalf: return ncclFloat16;
+        default:
+            TORCH_CHECK(false, "unsupported LoRA gradient dtype for EP all-reduce: ",
+                        tensor.scalar_type());
     }
+}
 
-    // Save groups for backward
-    ctx->group_ranges = groups;
+static void tp_broadcast_lora_parameter(
+    TrainingContext* ctx, at::Tensor& tensor
+) {
+    if (!ctx || ctx->tp_world_size <= 1 || !tensor.defined()) return;
+    TORCH_CHECK(ctx->tp_comm,
+        "LoRA TP communicator is not initialized for parameter broadcast");
+    TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous(),
+        "LoRA TP parameter broadcast requires a contiguous CUDA tensor");
+    const int dev = tensor.device().index();
+    cudaSetDevice(dev);
+    auto stream = c10::cuda::getCurrentCUDAStream(dev).stream();
+    auto err = ncclBroadcast(
+        tensor.data_ptr(), tensor.data_ptr(), tensor.numel(),
+        nccl_dtype_for(tensor), 0, ctx->tp_comm, stream);
+    TORCH_CHECK(err == ncclSuccess,
+        "NCCL LoRA parameter broadcast failed: ", ncclGetErrorString(err));
+}
 
-    for (auto& [start, end] : groups) {
-        if (offload && start < groups.back().first) {
-            ctx->group_inputs.push_back(
-                hidden.to(at::TensorOptions().dtype(hidden.scalar_type()).device(at::kCPU).pinned_memory(true))
-            );
-        } else {
-            ctx->group_inputs.push_back(hidden.clone());
-        }
+static void replica_broadcast_lora_parameter(
+    at::Tensor& tensor, ncclComm_t communicator, int world_size,
+    const char* axis
+) {
+    if (world_size <= 1 || !tensor.defined()) return;
+    TORCH_CHECK(communicator, "LoRA ", axis,
+        " communicator is not initialized for parameter broadcast");
+    TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous(),
+        "LoRA replica parameter broadcast requires a contiguous CUDA tensor");
+    const int dev = tensor.device().index();
+    cudaSetDevice(dev);
+    auto stream = c10::cuda::getCurrentCUDAStream(dev).stream();
+    auto err = ncclBroadcast(
+        tensor.data_ptr(), tensor.data_ptr(), tensor.numel(),
+        nccl_dtype_for(tensor), 0, communicator, stream);
+    TORCH_CHECK(err == ncclSuccess,
+        "NCCL LoRA ", axis, " parameter broadcast failed: ",
+        ncclGetErrorString(err));
+}
 
-        hidden = forward_layer_group(ctx, hidden, start, end);
-        // emptyCache: needed for seq>4096 (memory pressure), skip for small seq (async benefit).
-        if (hidden.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache();
+static void synchronize_adapter_replicated_lora_parameters(
+    TrainingContext* ctx, TrainingContext::LoRAAdapter& adapter);
+
+static void synchronize_fixed_replicated_lora_parameters(TrainingContext* ctx) {
+    if (!ctx) return;
+    for (int64_t layer = 0; layer < ctx->num_layers; ++layer) {
+        const int64_t offset = ctx->lora_layer_offset[layer];
+        const int64_t pairs = lora_pair_count(ctx->layer_configs[layer]);
+        for (int64_t pair = 0; pair < pairs; ++pair) {
+            const int64_t slot = offset + pair;
+            if (!legacy_lora_slot_active(ctx, slot)) continue;
+            const auto layout = lora_tp_layout(ctx, layer, pair);
+            if (layout == LoraTpLayout::ColumnParallel)
+                tp_broadcast_lora_parameter(ctx, ctx->lora_a[slot]);
+            else if (layout == LoraTpLayout::RowParallel)
+                tp_broadcast_lora_parameter(ctx, ctx->lora_b[slot]);
+            const bool grouped_expert =
+                lora_projection_table(ctx->layer_configs[layer])
+                    .entries[pair].grouped_expert;
+            if (!grouped_expert) {
+                replica_broadcast_lora_parameter(
+                    ctx->lora_a[slot], ctx->nccl_comm, ctx->ep_world_size, "EP");
+                replica_broadcast_lora_parameter(
+                    ctx->lora_b[slot], ctx->nccl_comm, ctx->ep_world_size, "EP");
+            }
+            replica_broadcast_lora_parameter(
+                ctx->lora_a[slot], ctx->dp_comm, ctx->dp_world_size, "DP");
+            replica_broadcast_lora_parameter(
+                ctx->lora_b[slot], ctx->dp_comm, ctx->dp_world_size, "DP");
+            replica_broadcast_lora_parameter(
+                ctx->lora_a[slot], ctx->cp_comm, ctx->cp_world_size, "CP");
+            replica_broadcast_lora_parameter(
+                ctx->lora_b[slot], ctx->cp_comm, ctx->cp_world_size, "CP");
+        }
+    }
+    for (auto& adapter : ctx->adapters)
+        synchronize_adapter_replicated_lora_parameters(ctx, adapter);
+}
 
-        if (getenv("QWEN36_DUMP_LAYERS")) {
-            auto h_f = hidden.to(at::kFloat);
-            fprintf(stderr, "[dump] ckpt group [%ld,%ld): mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
-                    (long)start, (long)end, h_f.mean().item<float>(), h_f.std().item<float>(),
-                    h_f[0][0][0].item<float>(), h_f[0][0][1].item<float>(), h_f[0][0][2].item<float>());
+static void synchronize_adapter_replicated_lora_parameters(
+    TrainingContext* ctx, TrainingContext::LoRAAdapter& adapter
+) {
+    if (!ctx) return;
+    if ((ctx->base_tp_attention || ctx->base_tp_mlp) &&
+        ctx->tp_world_size > 1 && !ctx->tp_comm)
+        return;  // qwen36_init_nccl synchronizes deferred adapters.
+    if (ctx->expert_parallel && ctx->ep_world_size > 1 && !ctx->nccl_comm)
+        return;
+    if (ctx->data_parallel && ctx->dp_world_size > 1 && !ctx->dp_comm)
+        return;
+    for (auto& [layer, pairs] : adapter.params) {
+        for (int64_t pair = 0; pair < static_cast<int64_t>(pairs.size()); ++pair) {
+            auto& [a, b] = pairs[pair];
+            if (!a.requires_grad() && !b.requires_grad()) continue;
+            const auto layout = lora_tp_layout(ctx, layer, pair);
+            if (layout == LoraTpLayout::ColumnParallel)
+                tp_broadcast_lora_parameter(ctx, a);
+            else if (layout == LoraTpLayout::RowParallel)
+                tp_broadcast_lora_parameter(ctx, b);
+            const bool grouped_expert =
+                lora_projection_table(ctx->layer_configs[layer])
+                    .entries[pair].grouped_expert;
+            if (!grouped_expert) {
+                replica_broadcast_lora_parameter(
+                    a, ctx->nccl_comm, ctx->ep_world_size, "EP");
+                replica_broadcast_lora_parameter(
+                    b, ctx->nccl_comm, ctx->ep_world_size, "EP");
+            }
+            replica_broadcast_lora_parameter(
+                a, ctx->dp_comm, ctx->dp_world_size, "DP");
+            replica_broadcast_lora_parameter(
+                b, ctx->dp_comm, ctx->dp_world_size, "DP");
+            replica_broadcast_lora_parameter(
+                a, ctx->cp_comm, ctx->cp_world_size, "CP");
+            replica_broadcast_lora_parameter(
+                b, ctx->cp_comm, ctx->cp_world_size, "CP");
         }
     }
+}
 
-    // Return hidden on GPU with requires_grad for CE backward
-    at::AutoGradMode restore(true);
-    hidden = hidden.set_requires_grad(true);
-    return hidden;
+static void tp_sum_replicated_lora_accumulator(
+    TrainingContext* ctx, at::Tensor& accumulator,
+    LoraTpLayout layout, bool is_a
+) {
+    const bool replicated =
+        (layout == LoraTpLayout::ColumnParallel && is_a) ||
+        (layout == LoraTpLayout::RowParallel && !is_a);
+    if (!replicated || !accumulator.defined() || ctx->tp_world_size <= 1) return;
+    TORCH_CHECK(ctx->tp_comm,
+        "LoRA TP communicator is not initialized for replicated gradient sum");
+    auto reduced = NcclAllReduceFunction::allreduce(
+        accumulator, ctx->tp_comm, ctx->tp_stream);
+    at::NoGradGuard guard;
+    accumulator.copy_(reduced);
 }
 
-// Manual sequential backward — recompute each group with grad, backprop, free.
-// Only 1 group's intermediate tensors exist at any time.
-static void manual_group_backward(
-    TrainingContext* ctx,
-    const at::Tensor& hidden_grad
+static void reduce_lora_accumulator(
+    TrainingContext* ctx, at::Tensor& accumulator, double scale,
+    bool reduce_ep, bool reduce_dp
 ) {
-    auto& groups = ctx->group_ranges;
-    int64_t num_groups = (int64_t)groups.size();
-    at::Tensor grad = hidden_grad;
-    at::AutoGradMode grad_mode(true);
+    if (!accumulator.defined()) return;
+    TORCH_CHECK(accumulator.scalar_type() == at::kFloat,
+        "LoRA DP gradient accumulator must be FP32");
+    auto contiguous = accumulator.contiguous();
+    if (scale != 1.0) {
+        contiguous = contiguous * scale;
+    }
+    auto reduce_axis = [&](ncclComm_t communicator, const char* axis) {
+        TORCH_CHECK(communicator, "LoRA ", axis,
+            " gradient all-reduce has no communicator");
+        auto reduced = at::empty_like(contiguous);
+        const int dev = contiguous.device().index();
+        cudaSetDevice(dev);
+        auto stream = c10::cuda::getCurrentCUDAStream(dev).stream();
+        auto err = ncclAllReduce(
+            contiguous.data_ptr(), reduced.data_ptr(), contiguous.numel(),
+            nccl_dtype_for(contiguous), ncclSum, communicator, stream);
+        TORCH_CHECK(err == ncclSuccess, "NCCL LoRA ", axis,
+            " gradient all-reduce failed: ", ncclGetErrorString(err));
+        contiguous = reduced;
+    };
+    if (reduce_ep) reduce_axis(ctx->nccl_comm, "EP");
+    if (reduce_dp) reduce_axis(ctx->dp_comm, "DP");
+    at::NoGradGuard guard;
+    accumulator.copy_(contiguous);
+}
 
-    ctx->lora_batch_valid = false;
-    ctx->lora_cache_valid = false;
-    // Rebuild lora_batch_cache under AutoGradMode(true) — this is critical:
-    // forward pass built it in no-grad mode, so a_stack/b_stack didn't track
-    // gradients. Here we rebuild so at::stack(A_i) creates a differentiable
-    // graph: grad → bmm → stack → A_i (leaf). Without this, grads are undefined.
-    prepare_lora_batch(ctx);
+static void normalize_lora_accumulator_numerator(
+    TrainingContext* ctx, at::Tensor& accumulator,
+    const at::Tensor& global_weight, bool reduce_ep, bool reduce_dp
+) {
+    if (!accumulator.defined()) return;
+    TORCH_CHECK(accumulator.scalar_type() == at::kFloat,
+        "LoRA DP gradient accumulator must be FP32");
+    TORCH_CHECK(global_weight.numel() == 1,
+        "per-adapter LoRA global token weight must be scalar");
+    reduce_lora_accumulator(
+        ctx, accumulator, 1.0, reduce_ep, reduce_dp);
+    at::NoGradGuard guard;
+    accumulator.div_(global_weight.clamp_min(1.0));
+}
 
-    for (int64_t g = num_groups - 1; g >= 0; g--) {
-        int64_t start = groups[g].first;
-        int64_t end = groups[g].second;
+static void reduce_lora_accumulator_weighted(
+    TrainingContext* ctx, at::Tensor& accumulator,
+    const at::Tensor& local_weight, const at::Tensor& global_weight,
+    bool reduce_ep, bool reduce_dp
+) {
+    if (!accumulator.defined()) return;
+    TORCH_CHECK(accumulator.scalar_type() == at::kFloat,
+        "LoRA DP gradient accumulator must be FP32");
+    TORCH_CHECK(local_weight.numel() == 1 && global_weight.numel() == 1,
+        "per-adapter LoRA token weights must be scalar");
+    at::NoGradGuard guard;
+    accumulator.mul_(local_weight);
+    reduce_lora_accumulator(
+        ctx, accumulator, 1.0, reduce_ep, reduce_dp);
+    accumulator.div_(global_weight.clamp_min(1.0));
+}
 
-        // Restore input from saved (CPU if offloaded)
-        auto input = ctx->group_inputs[g].to(hidden_grad.device()).detach().set_requires_grad(true);
+struct GroupedLoraGradientSyncPlan {
+    static constexpr uint8_t kReduceEp = 1 << 0;
+    static constexpr uint8_t kReduceDp = 1 << 1;
+    static constexpr uint8_t kReduceTp = 1 << 2;
+    static constexpr uint8_t kReduceCp = 1 << 3;
+
+    struct Entry {
+        at::Tensor* accumulator = nullptr;
+        at::Tensor local_weight;
+        at::Tensor global_weight;
+        at::Tensor work;
+        double pre_scale = 1.0;
+        double post_scale = 1.0;
+        bool reduce_ep = false;
+        bool reduce_dp = false;
+        bool reduce_tp = false;
+        bool reduce_cp = false;
+    };
 
-        // Recompute forward with grad for this group only
-        auto output = forward_layer_group(ctx, input, start, end);
+    struct Bucket {
+        uint8_t communication_mask = 0;
+        std::vector<Entry*> entries;
+        at::Tensor storage;
+    };
 
-        // Backprop through this group using grad() instead of backward().
-        // grad() only computes gradients for specified inputs — faster than
-        // backward() which traverses all leaf nodes.
-        // LoRA params are shared across groups, so we accumulate their gradients.
-        std::vector<at::Tensor> grad_inputs = {input};
+    std::vector<Entry> entries;
+    std::vector<Bucket> buckets;
+    // Dynamic CP keeps LoRA parameters replicated while each rank owns a
+    // different sequence/head contribution. Sum those gradients here; the
+    // full-batch token denominator remains replicated. Fixed LoRA keeps its
+    // legacy CP reduction after this plan.
+    bool reduce_cp = false;
+
+    static bool tp_replicated(LoraTpLayout layout, bool is_a) {
+        return (layout == LoraTpLayout::ColumnParallel && is_a) ||
+            (layout == LoraTpLayout::RowParallel && !is_a);
+    }
 
-        if (ctx->lora_batch_valid) {
-            // Multi-LoRA: collect A/B from ctx->adapters
-            for (int64_t l = start; l < end; l++) {
-                int64_t lora_count = (ctx->layer_configs[l].layer_type == 0) ? 4 : 3;
-                for (auto& adapter : ctx->adapters) {
-                    auto it = adapter.params.find(l);
-                    if (it == adapter.params.end()) continue;
-                    for (int64_t k = 0; k < lora_count && k < (int64_t)it->second.size(); k++) {
-                        auto& [a, b] = it->second[k];
-                        grad_inputs.push_back(a);
-                        grad_inputs.push_back(b);
-                    }
-                }
+    void add(
+        TrainingContext* ctx,
+        at::Tensor& accumulator,
+        double pre_scale,
+        const at::Tensor& local_weight,
+        const at::Tensor& global_weight,
+        double post_scale,
+        bool reduce_ep,
+        bool reduce_dp,
+        LoraTpLayout tp_layout,
+        bool is_a
+    ) {
+        if (!accumulator.defined()) return;
+        Entry entry;
+        entry.accumulator = &accumulator;
+        entry.local_weight = local_weight;
+        entry.global_weight = global_weight;
+        entry.pre_scale = pre_scale;
+        entry.post_scale = post_scale;
+        entry.reduce_ep = reduce_ep;
+        entry.reduce_dp = reduce_dp;
+        entry.reduce_tp = ctx->tp_world_size > 1 &&
+            tp_replicated(tp_layout, is_a);
+        entry.reduce_cp = reduce_cp;
+        entries.push_back(std::move(entry));
+    }
+
+    static uint8_t communication_mask(const Entry& entry) {
+        return (entry.reduce_ep ? kReduceEp : 0) |
+            (entry.reduce_dp ? kReduceDp : 0) |
+            (entry.reduce_tp ? kReduceTp : 0) |
+            (entry.reduce_cp ? kReduceCp : 0);
+    }
+
+    void prepare(TrainingContext* ctx, bool packed_sync) {
+        buckets.clear();
+        for (auto& entry : entries) {
+            TORCH_CHECK(entry.accumulator && entry.accumulator->is_cuda() &&
+                    entry.accumulator->scalar_type() == at::kFloat &&
+                    entry.accumulator->device().index() == ctx->cuda_device,
+                "grouped LoRA gradient sync requires CUDA FP32 accumulators "
+                "on the context device");
+            TORCH_CHECK(std::isfinite(entry.pre_scale) &&
+                    std::isfinite(entry.post_scale),
+                "grouped LoRA gradient sync scale must be finite");
+            if (entry.local_weight.defined()) {
+                TORCH_CHECK(entry.local_weight.is_cuda() &&
+                        entry.local_weight.numel() == 1 &&
+                        entry.local_weight.device() == entry.accumulator->device(),
+                    "grouped LoRA local token weight must be a CUDA scalar");
             }
-        } else {
-            // Legacy single-LoRA
-            for (int64_t l = start; l < end; l++) {
-                int64_t lora_count = (ctx->layer_configs[l].layer_type == 0) ? 4 : 3;
-                int64_t la_offset = ctx->lora_layer_offset[l];
-                bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size();
-                if (has_lora) {
-                    for (int64_t k = 0; k < lora_count; k++) {
-                        grad_inputs.push_back(ctx->lora_a[la_offset + k]);
-                        grad_inputs.push_back(ctx->lora_b[la_offset + k]);
-                    }
-                }
+            if (entry.global_weight.defined()) {
+                TORCH_CHECK(entry.global_weight.is_cuda() &&
+                        entry.global_weight.numel() == 1 &&
+                        entry.global_weight.device() == entry.accumulator->device(),
+                    "grouped LoRA global token weight must be a CUDA scalar");
+            }
+            TORCH_CHECK(entry.accumulator->is_contiguous(),
+                "grouped LoRA gradient sync requires contiguous accumulators");
+            entry.work = *entry.accumulator;
+            const uint8_t mask = communication_mask(entry);
+            if (!packed_sync || mask == 0) continue;
+            auto bucket = std::find_if(
+                buckets.begin(), buckets.end(),
+                [mask](const Bucket& candidate) {
+                    return candidate.communication_mask == mask;
+                });
+            if (bucket == buckets.end()) {
+                buckets.push_back(Bucket{});
+                bucket = std::prev(buckets.end());
+                bucket->communication_mask = mask;
             }
+            bucket->entries.push_back(&entry);
         }
+        for (auto& bucket : buckets) {
+            if (bucket.entries.size() < 2) continue;
+            std::vector<at::Tensor> flattened;
+            flattened.reserve(bucket.entries.size());
+            for (const auto* entry : bucket.entries)
+                flattened.push_back(entry->work.reshape({-1}));
+            bucket.storage = at::cat(flattened, 0);
+            int64_t offset = 0;
+            for (auto* entry : bucket.entries) {
+                const int64_t elements = entry->accumulator->numel();
+                entry->work = bucket.storage.narrow(0, offset, elements)
+                    .view(entry->accumulator->sizes());
+                offset += elements;
+            }
+        }
+    }
 
-        auto grads = torch::autograd::grad(
-            {output}, grad_inputs, {grad},
-            /*retain_graph=*/false, /*create_graph=*/false,
-            /*allow_unused=*/true
-        );
+    template <typename Predicate>
+    static void run_group(
+        std::vector<Entry>& entries,
+        ncclComm_t communicator,
+        cudaStream_t stream,
+        const char* axis,
+        Predicate predicate
+    ) {
+        const bool any = std::any_of(
+            entries.begin(), entries.end(), predicate);
+        if (!any) return;
+        TORCH_CHECK(communicator, "grouped LoRA ", axis,
+            " gradient sync has no communicator");
+        const auto start_error = ncclGroupStart();
+        TORCH_CHECK(start_error == ncclSuccess,
+            "grouped LoRA ", axis, " ncclGroupStart failed: ",
+            ncclGetErrorString(start_error));
+        ncclResult_t first_error = ncclSuccess;
+        for (auto& entry : entries) {
+            if (!predicate(entry)) continue;
+            const auto error = ncclAllReduce(
+                entry.work.data_ptr(), entry.work.data_ptr(), entry.work.numel(),
+                nccl_dtype_for(entry.work), ncclSum, communicator, stream);
+            if (first_error == ncclSuccess && error != ncclSuccess)
+                first_error = error;
+        }
+        const auto end_error = ncclGroupEnd();
+        TORCH_CHECK(first_error == ncclSuccess,
+            "grouped LoRA ", axis, " gradient all-reduce failed: ",
+            ncclGetErrorString(first_error));
+        TORCH_CHECK(end_error == ncclSuccess,
+            "grouped LoRA ", axis, " ncclGroupEnd failed: ",
+            ncclGetErrorString(end_error));
+    }
+
+    static void run_bucket_group(
+        std::vector<Bucket>& buckets,
+        ncclComm_t communicator,
+        cudaStream_t stream,
+        const char* axis,
+        uint8_t axis_mask
+    ) {
+        const bool any = std::any_of(
+            buckets.begin(), buckets.end(),
+            [axis_mask](const Bucket& bucket) {
+                return bucket.entries.size() >= 2 &&
+                    (bucket.communication_mask & axis_mask) != 0;
+            });
+        if (!any) return;
+        TORCH_CHECK(communicator, "packed LoRA ", axis,
+            " gradient sync has no communicator");
+        const auto start_error = ncclGroupStart();
+        TORCH_CHECK(start_error == ncclSuccess,
+            "packed LoRA ", axis, " ncclGroupStart failed: ",
+            ncclGetErrorString(start_error));
+        ncclResult_t first_error = ncclSuccess;
+        for (auto& bucket : buckets) {
+            if (bucket.entries.size() < 2 ||
+                (bucket.communication_mask & axis_mask) == 0)
+                continue;
+            const auto error = ncclAllReduce(
+                bucket.storage.data_ptr(), bucket.storage.data_ptr(),
+                bucket.storage.numel(), nccl_dtype_for(bucket.storage),
+                ncclSum, communicator, stream);
+            if (first_error == ncclSuccess && error != ncclSuccess)
+                first_error = error;
+        }
+        const auto end_error = ncclGroupEnd();
+        TORCH_CHECK(first_error == ncclSuccess,
+            "packed LoRA ", axis, " gradient all-reduce failed: ",
+            ncclGetErrorString(first_error));
+        TORCH_CHECK(end_error == ncclSuccess,
+            "packed LoRA ", axis, " ncclGroupEnd failed: ",
+            ncclGetErrorString(end_error));
+    }
 
+    void execute(TrainingContext* ctx) {
+        const bool packed_sync = env_enabled(
+            "QWEN36_PACKED_LORA_SYNC", true);
+        std::exception_ptr preparation_error;
+        try {
+            prepare(ctx, packed_sync);
+        } catch (...) {
+            preparation_error = std::current_exception();
+        }
+        const bool local_ready = !preparation_error;
+        const bool globally_ready = adapter_collective_all_succeeded(
+            ctx, local_ready);
+        if (!local_ready) std::rethrow_exception(preparation_error);
+        TORCH_CHECK(globally_ready,
+            "grouped LoRA gradient sync preparation failed on another rank");
+
+        at::NoGradGuard guard;
+        for (auto& entry : entries) {
+            if (entry.local_weight.defined())
+                entry.work.mul_(entry.local_weight);
+            if (entry.pre_scale != 1.0)
+                entry.work.mul_(entry.pre_scale);
+        }
 
-        // Manually accumulate LoRA param gradients
-        if (ctx->lora_batch_valid) {
-            // Multi-LoRA: accumulate into ctx->adapters
-            int64_t gi = 1;  // skip input grad (index 0)
-            for (int64_t l = start; l < end; l++) {
-                int64_t lora_count = (ctx->layer_configs[l].layer_type == 0) ? 4 : 3;
-                for (auto& adapter : ctx->adapters) {
-                    auto it = adapter.params.find(l);
-                    if (it == adapter.params.end()) continue;
-                    for (int64_t k = 0; k < lora_count && k < (int64_t)it->second.size(); k++) {
-                        auto& [a, b] = it->second[k];
-                        if (gi < (int64_t)grads.size() && grads[gi].defined()) {
-                            if (a.grad().defined()) a.grad().add_(grads[gi]);
-                            else a.mutable_grad() = grads[gi].clone();
-                        }
-                        gi++;
-                        if (gi < (int64_t)grads.size() && grads[gi].defined()) {
-                            if (b.grad().defined()) b.grad().add_(grads[gi]);
-                            else b.mutable_grad() = grads[gi].clone();
-                        }
-                        gi++;
-                    }
-                }
-            }
+        const auto stream = c10::cuda::getCurrentCUDAStream(
+            ctx->cuda_device).stream();
+        if (packed_sync) {
+            run_bucket_group(
+                buckets, ctx->cp_comm, stream, "CP", kReduceCp);
+            run_bucket_group(
+                buckets, ctx->nccl_comm, stream, "EP", kReduceEp);
+            run_bucket_group(
+                buckets, ctx->dp_comm, stream, "DP", kReduceDp);
+            run_group(entries, ctx->cp_comm, stream, "CP",
+                [](const Entry& entry) {
+                    return entry.reduce_cp &&
+                        entry.work.data_ptr() ==
+                            entry.accumulator->data_ptr();
+                });
+            run_group(entries, ctx->nccl_comm, stream, "EP",
+                [](const Entry& entry) {
+                    return entry.reduce_ep &&
+                        entry.work.data_ptr() ==
+                            entry.accumulator->data_ptr();
+                });
+            run_group(entries, ctx->dp_comm, stream, "DP",
+                [](const Entry& entry) {
+                    return entry.reduce_dp &&
+                        entry.work.data_ptr() ==
+                            entry.accumulator->data_ptr();
+                });
         } else {
-            // Legacy single-LoRA
-            int64_t gi = 1;  // skip input grad (index 0)
-            for (int64_t l = start; l < end; l++) {
-                int64_t lora_count = (ctx->layer_configs[l].layer_type == 0) ? 4 : 3;
-                int64_t la_offset = ctx->lora_layer_offset[l];
-                bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size();
-                if (has_lora) {
-                    for (int64_t k = 0; k < lora_count; k++) {
-                        if (grads[gi].defined()) {
-                            auto& pa = ctx->lora_a[la_offset + k];
-                            if (pa.grad().defined()) pa.grad().add_(grads[gi]);
-                            else pa.mutable_grad() = grads[gi].clone();
+            run_group(entries, ctx->cp_comm, stream, "CP",
+                [](const Entry& entry) { return entry.reduce_cp; });
+            run_group(entries, ctx->nccl_comm, stream, "EP",
+                [](const Entry& entry) { return entry.reduce_ep; });
+            run_group(entries, ctx->dp_comm, stream, "DP",
+                [](const Entry& entry) { return entry.reduce_dp; });
+        }
+
+        for (auto& entry : entries) {
+            if (entry.global_weight.defined())
+                entry.work.div_(entry.global_weight);
+            if (entry.post_scale != 1.0)
+                entry.work.mul_(entry.post_scale);
+        }
+
+        if (packed_sync) {
+            run_bucket_group(
+                buckets, ctx->tp_comm, stream, "TP", kReduceTp);
+            run_group(entries, ctx->tp_comm, stream, "TP",
+                [](const Entry& entry) {
+                    return entry.reduce_tp &&
+                        entry.work.data_ptr() ==
+                            entry.accumulator->data_ptr();
+                });
+            for (auto& entry : entries) {
+                if (entry.work.data_ptr() != entry.accumulator->data_ptr())
+                    entry.accumulator->copy_(entry.work);
+            }
+        } else {
+            run_group(entries, ctx->tp_comm, stream, "TP",
+                [](const Entry& entry) { return entry.reduce_tp; });
+        }
+    }
+};
+
+// Fixed-LoRA gradients are accumulated as token-weighted numerators. Replicated
+// DP sums all replicated parameters and divides by the global token count;
+// legacy EP keeps its replicated batch/local expert semantics. Sharded A2A
+// sums non-expert parameters across source ranks, while expert parameters have
+// already received all source numerators through the inverse A2A and therefore
+// only divide by the global token count.
+static bool replica_token_weights_match(
+    TrainingContext* ctx, const at::Tensor& weights,
+    ncclComm_t communicator, int32_t world_size, const char* axis
+) {
+    if (!communicator || world_size <= 1) return true;
+    TORCH_CHECK(weights.is_cuda() && weights.scalar_type() == at::kFloat &&
+            weights.is_contiguous(),
+        "LoRA replica token weights must be contiguous CUDA FP32");
+    auto minimum = weights.clone();
+    auto maximum = weights.clone();
+    auto stream = c10::cuda::getCurrentCUDAStream(
+        weights.device().index()).stream();
+    const auto minimum_error = ncclAllReduce(
+        minimum.data_ptr<float>(), minimum.data_ptr<float>(), minimum.numel(),
+        ncclFloat, ncclMin, communicator, stream);
+    TORCH_CHECK(minimum_error == ncclSuccess,
+        "NCCL LoRA ", axis,
+        " token-count minimum validation failed: ",
+        ncclGetErrorString(minimum_error));
+    const auto maximum_error = ncclAllReduce(
+        maximum.data_ptr<float>(), maximum.data_ptr<float>(), maximum.numel(),
+        ncclFloat, ncclMax, communicator, stream);
+    TORCH_CHECK(maximum_error == ncclSuccess,
+        "NCCL LoRA ", axis,
+        " token-count maximum validation failed: ",
+        ncclGetErrorString(maximum_error));
+    return (minimum == maximum).all().item<bool>();
+}
+
+static void cp_sum_fixed_lora_accumulators(TrainingContext* ctx) {
+    if (!context_parallel_enabled(ctx)) return;
+    TORCH_CHECK(ctx->cp_comm && ctx->adapters.empty(),
+        "CP gradient synchronization requires fixed LoRA and a CP communicator");
+    auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+    if (ctx->fixed_grad_slab.defined()) {
+        const auto error = ncclAllReduce(
+            ctx->fixed_grad_slab.data_ptr<float>(),
+            ctx->fixed_grad_slab.data_ptr<float>(),
+            ctx->fixed_grad_slab.numel(), ncclFloat, ncclSum,
+            ctx->cp_comm, stream);
+        TORCH_CHECK(error == ncclSuccess,
+            "CP packed LoRA gradient all-reduce failed: ",
+            ncclGetErrorString(error));
+        return;
+    }
+    for (size_t index = 0; index < ctx->lora_a.size(); ++index) {
+        if (!ctx->lora_active[index]) continue;
+        for (auto* accumulator : {
+                &ctx->grad_accum_a[index], &ctx->grad_accum_b[index]}) {
+            TORCH_CHECK(accumulator->defined() &&
+                    accumulator->scalar_type() == at::kFloat &&
+                    accumulator->is_cuda() && accumulator->is_contiguous(),
+                "CP LoRA gradient accumulator must be contiguous CUDA FP32");
+        }
+    }
+    TORCH_CHECK(ncclGroupStart() == ncclSuccess,
+        "CP LoRA gradient group start failed");
+    ncclResult_t first_error = ncclSuccess;
+    for (size_t index = 0; index < ctx->lora_a.size(); ++index) {
+        if (!ctx->lora_active[index]) continue;
+        for (auto* accumulator : {
+                &ctx->grad_accum_a[index], &ctx->grad_accum_b[index]}) {
+            const auto error = ncclAllReduce(
+                accumulator->data_ptr<float>(), accumulator->data_ptr<float>(),
+                accumulator->numel(), ncclFloat, ncclSum,
+                ctx->cp_comm, stream);
+            if (first_error == ncclSuccess && error != ncclSuccess)
+                first_error = error;
+        }
+    }
+    const auto group_error = ncclGroupEnd();
+    TORCH_CHECK(first_error == ncclSuccess,
+        "CP LoRA gradient all-reduce failed: ",
+        ncclGetErrorString(first_error));
+    TORCH_CHECK(group_error == ncclSuccess,
+        "CP LoRA gradient group end failed: ",
+        ncclGetErrorString(group_error));
+}
+
+static bool synchronize_lora_gradients(
+    TrainingContext* ctx, const at::Tensor& target_mask,
+    double accumulated_token_weight = 0.0,
+    const at::Tensor* per_adapter_token_counts = nullptr,
+    std::vector<uint8_t>* adapter_has_global_tokens = nullptr,
+    bool adapter_token_counts_prevalidated = false,
+    bool accumulators_are_numerators = false
+) {
+    const bool sharded_a2a = ctx->expert_parallel && ctx->nccl_comm &&
+        env_enabled("QWEN36_EP_A2A_SHARDED");
+    TORCH_CHECK(!sharded_a2a || env_enabled("QWEN36_EP_A2A"),
+        "QWEN36_EP_A2A_SHARDED=1 requires QWEN36_EP_A2A=1");
+    const bool dp_allreduce = ctx->data_parallel && ctx->dp_comm;
+    const bool per_adapter_weighting = per_adapter_token_counts &&
+        per_adapter_token_counts->defined();
+    const bool dynamic_cp = per_adapter_weighting &&
+        context_parallel_enabled(ctx);
+    const bool normalization_allreduce = dp_allreduce || sharded_a2a;
+    const double replicated_a2a_expert_scale =
+        ctx->expert_parallel && env_enabled("QWEN36_EP_A2A") && !sharded_a2a
+        ? 1.0 / static_cast<double>(ctx->ep_world_size)
+        : 1.0;
+    auto sum_replica_axes = [&](at::Tensor value, bool reduce_ep) {
+        auto reduce_axis = [&](ncclComm_t communicator, const char* axis) {
+            auto reduced = at::empty_like(value);
+            auto stream = c10::cuda::getCurrentCUDAStream(
+                value.device().index()).stream();
+            auto err = ncclAllReduce(
+                value.data_ptr(), reduced.data_ptr(), value.numel(),
+                nccl_dtype_for(value), ncclSum, communicator, stream);
+            TORCH_CHECK(err == ncclSuccess, "NCCL ", axis,
+                " replica all-reduce failed: ", ncclGetErrorString(err));
+            value = reduced;
+        };
+        if (reduce_ep) reduce_axis(ctx->nccl_comm, "EP");
+        if (dp_allreduce) reduce_axis(ctx->dp_comm, "DP");
+        return value;
+    };
+    const bool grouped_sync = env_enabled("QWEN36_GROUPED_LORA_SYNC", true);
+    TORCH_CHECK(!dynamic_cp || grouped_sync,
+        "dynamic LoRA with context parallelism requires grouped gradient sync");
+    GroupedLoraGradientSyncPlan grouped_plan;
+    grouped_plan.reduce_cp = dynamic_cp;
+    at::Tensor local_adapter_weights;
+    at::Tensor global_adapter_weights;
+    double scale = 1.0;
+    if (per_adapter_weighting) {
+        TORCH_CHECK(per_adapter_token_counts->dim() == 1 &&
+                    per_adapter_token_counts->size(0) ==
+                        static_cast<int64_t>(ctx->adapters.size()),
+            "dynamic LoRA token-count vector must match adapter registry");
+        local_adapter_weights = per_adapter_token_counts->to(at::kFloat)
+            .contiguous();
+        if (!adapter_token_counts_prevalidated) {
+            TORCH_CHECK(at::logical_and(
+                    at::isfinite(local_adapter_weights),
+                    local_adapter_weights >= 0).all().item<bool>(),
+                "dynamic LoRA token counts must be finite and non-negative");
+        }
+        bool replica_weights_match = replica_token_weights_match(
+            ctx, local_adapter_weights, ctx->tp_comm,
+            ctx->tp_world_size, "TP");
+        replica_weights_match = replica_weights_match &&
+            replica_token_weights_match(
+                ctx, local_adapter_weights, ctx->cp_comm,
+                ctx->cp_world_size, "CP");
+        if (ctx->expert_parallel && !sharded_a2a) {
+            const bool ep_weights_match = replica_token_weights_match(
+                ctx, local_adapter_weights, ctx->nccl_comm,
+                ctx->ep_world_size, "replicated EP");
+            replica_weights_match = replica_weights_match && ep_weights_match;
+        }
+        TORCH_CHECK(adapter_collective_all_succeeded(
+                ctx, replica_weights_match) && replica_weights_match,
+            "dynamic LoRA token counts differ across TP or replicated EP "
+            "ranks; replicas must use identical target masks");
+        global_adapter_weights = normalization_allreduce
+            ? sum_replica_axes(local_adapter_weights, sharded_a2a)
+            : local_adapter_weights;
+        if (adapter_has_global_tokens) {
+            adapter_has_global_tokens->assign(ctx->adapters.size(), 0);
+            auto global_cpu = global_adapter_weights.to(
+                at::TensorOptions().device(at::kCPU));
+            const auto* counts = global_cpu.data_ptr<float>();
+            for (size_t i = 0; i < ctx->adapters.size(); ++i) {
+                (*adapter_has_global_tokens)[i] = counts[i] > 0.0f ? 1 : 0;
+            }
+        }
+    } else {
+        auto local = at::full({1}, accumulated_token_weight,
+            at::TensorOptions().dtype(at::kFloat).device(target_mask.device()));
+        bool replica_weights_match = replica_token_weights_match(
+            ctx, local, ctx->tp_comm, ctx->tp_world_size, "TP");
+        replica_weights_match = replica_weights_match &&
+            replica_token_weights_match(
+                ctx, local, ctx->cp_comm, ctx->cp_world_size, "CP");
+        if (ctx->expert_parallel && !sharded_a2a) {
+            const bool ep_weights_match = replica_token_weights_match(
+                ctx, local, ctx->nccl_comm,
+                ctx->ep_world_size, "replicated EP");
+            replica_weights_match = replica_weights_match && ep_weights_match;
+        }
+        TORCH_CHECK(adapter_collective_all_succeeded(
+                ctx, replica_weights_match) && replica_weights_match,
+            "fixed LoRA token weights differ across TP, CP, or replicated EP "
+            "ranks; replicas must use identical accumulation windows");
+        auto global = normalization_allreduce
+            ? sum_replica_axes(local, sharded_a2a)
+            : local;
+        const double global_weight = global.item<double>();
+        if (global_weight <= 0.0) return false;
+        scale = 1.0 / global_weight;
+    }
+    for (size_t adapter_index = 0;
+         adapter_index < ctx->adapters.size(); ++adapter_index) {
+        auto& adapter = ctx->adapters[adapter_index];
+        for (auto& [layer_idx, pairs] : adapter.params) {
+            auto table = lora_projection_table(ctx->layer_configs[layer_idx]);
+            for (int64_t pair = 0; pair < (int64_t)pairs.size(); ++pair) {
+                auto accum_it = adapter.grad_accum.find(layer_idx);
+                TORCH_CHECK(accum_it != adapter.grad_accum.end() &&
+                            pair < (int64_t)accum_it->second.size(),
+                    "dynamic LoRA gradient accumulator layout mismatch");
+                if (per_adapter_weighting) {
+                    auto local_weight = local_adapter_weights.index(
+                        {static_cast<int64_t>(adapter_index)});
+                    auto global_weight = global_adapter_weights.index(
+                        {static_cast<int64_t>(adapter_index)});
+                    TORCH_CHECK(adapter_has_global_tokens &&
+                            adapter_has_global_tokens->size() == ctx->adapters.size(),
+                        "dynamic LoRA global-token activity is unavailable");
+                    if (!(*adapter_has_global_tokens)[adapter_index]) {
+                        if (accum_it->second[pair][0].defined())
+                            accum_it->second[pair][0].zero_();
+                        if (accum_it->second[pair][1].defined())
+                            accum_it->second[pair][1].zero_();
+                        continue;
+                    }
+                    const bool grouped_expert = table.entries[pair].grouped_expert;
+                    const auto layout = lora_tp_layout(ctx, layer_idx, pair);
+                    if (sharded_a2a) {
+                        // Sharded A2A restores each source row to a numerator
+                        // before backward. Expert owners already receive all
+                        // source numerators; replicated projections still need
+                        // one all-reduce before division by the global count.
+                        if (grouped_sync) {
+                            grouped_plan.add(
+                                ctx, accum_it->second[pair][0], 1.0,
+                                at::Tensor(), global_weight, 1.0,
+                                !grouped_expert, dp_allreduce, layout, true);
+                            grouped_plan.add(
+                                ctx, accum_it->second[pair][1], 1.0,
+                                at::Tensor(), global_weight, 1.0,
+                                !grouped_expert, dp_allreduce, layout, false);
+                        } else {
+                            normalize_lora_accumulator_numerator(
+                                ctx, accum_it->second[pair][0], global_weight,
+                                !grouped_expert, dp_allreduce);
+                            normalize_lora_accumulator_numerator(
+                                ctx, accum_it->second[pair][1], global_weight,
+                                !grouped_expert, dp_allreduce);
                         }
-                        gi++;
-                        if (grads[gi].defined()) {
-                            auto& pb = ctx->lora_b[la_offset + k];
-                            if (pb.grad().defined()) pb.grad().add_(grads[gi]);
-                            else pb.mutable_grad() = grads[gi].clone();
+                    } else {
+                        if (grouped_sync) {
+                            const double expert_scale = grouped_expert
+                                ? replicated_a2a_expert_scale
+                                : 1.0;
+                            grouped_plan.add(
+                                ctx, accum_it->second[pair][0], 1.0,
+                                accumulators_are_numerators ? at::Tensor() : local_weight,
+                                global_weight, expert_scale,
+                                false, dp_allreduce, layout, true);
+                            grouped_plan.add(
+                                ctx, accum_it->second[pair][1], 1.0,
+                                accumulators_are_numerators ? at::Tensor() : local_weight,
+                                global_weight, expert_scale,
+                                false, dp_allreduce, layout, false);
+                        } else {
+                            if (accumulators_are_numerators) {
+                                normalize_lora_accumulator_numerator(
+                                    ctx, accum_it->second[pair][0], global_weight,
+                                    false, dp_allreduce);
+                                normalize_lora_accumulator_numerator(
+                                    ctx, accum_it->second[pair][1], global_weight,
+                                    false, dp_allreduce);
+                            } else {
+                                reduce_lora_accumulator_weighted(
+                                    ctx, accum_it->second[pair][0], local_weight,
+                                    global_weight, false, dp_allreduce);
+                                reduce_lora_accumulator_weighted(
+                                    ctx, accum_it->second[pair][1], local_weight,
+                                    global_weight, false, dp_allreduce);
+                            }
+                        }
+                        if (!grouped_sync && grouped_expert &&
+                            replicated_a2a_expert_scale != 1.0) {
+                            at::NoGradGuard guard;
+                            accum_it->second[pair][0].mul_(
+                                replicated_a2a_expert_scale);
+                            accum_it->second[pair][1].mul_(
+                                replicated_a2a_expert_scale);
+                        }
+                    }
+                    continue;
+                }
+                // Routed-expert tensors are sharded only in EP. Pure DP has
+                // replicated experts and therefore uses the same reduction as
+                // shared projections; the legacy accumulation path preserves
+                // local normalization when no DP communicator is present.
+                if (table.entries[pair].grouped_expert) {
+                    if (accumulated_token_weight > 0.0) {
+                        const auto layout = lora_tp_layout(ctx, layer_idx, pair);
+                        if (grouped_sync) {
+                            grouped_plan.add(
+                                ctx, accum_it->second[pair][0], scale,
+                                at::Tensor(), at::Tensor(), 1.0,
+                                false, dp_allreduce, layout, true);
+                            grouped_plan.add(
+                                ctx, accum_it->second[pair][1], scale,
+                                at::Tensor(), at::Tensor(), 1.0,
+                                false, dp_allreduce, layout, false);
+                        } else {
+                            reduce_lora_accumulator(
+                                ctx, accum_it->second[pair][0], scale,
+                                false, dp_allreduce);
+                            reduce_lora_accumulator(
+                                ctx, accum_it->second[pair][1], scale,
+                                false, dp_allreduce);
                         }
-                        gi++;
                     }
+                    continue;
+                }
+                const auto layout = lora_tp_layout(ctx, layer_idx, pair);
+                if (grouped_sync) {
+                    grouped_plan.add(
+                        ctx, accum_it->second[pair][0], scale,
+                        at::Tensor(), at::Tensor(), 1.0,
+                        sharded_a2a, dp_allreduce, layout, true);
+                    grouped_plan.add(
+                        ctx, accum_it->second[pair][1], scale,
+                        at::Tensor(), at::Tensor(), 1.0,
+                        sharded_a2a, dp_allreduce, layout, false);
+                } else {
+                    reduce_lora_accumulator(
+                        ctx, accum_it->second[pair][0], scale,
+                        sharded_a2a, dp_allreduce);
+                    reduce_lora_accumulator(
+                        ctx, accum_it->second[pair][1], scale,
+                        sharded_a2a, dp_allreduce);
+                }
+            }
+        }
+    }
+    // Projection-aware TP keeps one LoRA factor replicated. Its gradient is
+    // the sum of disjoint output-head (column) or input-column (row)
+    // contributions and is synchronized once at the optimizer boundary.
+    if (!grouped_sync) {
+        for (auto& adapter : ctx->adapters) {
+            for (auto& [layer_idx, pairs] : adapter.grad_accum) {
+                for (int64_t pair = 0;
+                     pair < static_cast<int64_t>(pairs.size()); ++pair) {
+                    const auto layout = lora_tp_layout(ctx, layer_idx, pair);
+                    tp_sum_replicated_lora_accumulator(
+                        ctx, pairs[pair][0], layout, true);
+                    tp_sum_replicated_lora_accumulator(
+                        ctx, pairs[pair][1], layout, false);
+                }
+            }
+        }
+    }
+    if (per_adapter_weighting) {
+        if (grouped_sync) grouped_plan.execute(ctx);
+        return adapter_has_global_tokens && std::any_of(
+            adapter_has_global_tokens->begin(),
+            adapter_has_global_tokens->end(),
+            [](uint8_t active) { return active != 0; });
+    }
+    // The replicated-source A2A path sends an identical token batch from every
+    // EP rank. Average only its sharded expert parameter gradients here;
+    // scaling the combine backward would also under-scale the activation
+    // gradient returned to each source graph. Sharded A2A is handled above via
+    // global token-count weighting and does not enter this branch.
+    for (int64_t layer = 0; layer < ctx->num_layers; ++layer) {
+        auto table = lora_projection_table(ctx->layer_configs[layer]);
+        int64_t offset = ctx->lora_layer_offset[layer];
+        for (int64_t pair = 0; pair < table.count; ++pair) {
+            // Routed expert LoRA is local-only for EP, while pure DP owns the
+            // complete replicated expert tensor and must all-reduce it.
+            if (table.entries[pair].grouped_expert) {
+                const auto layout = lora_tp_layout(ctx, layer, pair);
+                if (grouped_sync) {
+                    grouped_plan.add(
+                        ctx, ctx->grad_accum_a[offset + pair],
+                        scale * replicated_a2a_expert_scale,
+                        at::Tensor(), at::Tensor(), 1.0,
+                        false, dp_allreduce, layout, true);
+                    grouped_plan.add(
+                        ctx, ctx->grad_accum_b[offset + pair],
+                        scale * replicated_a2a_expert_scale,
+                        at::Tensor(), at::Tensor(), 1.0,
+                        false, dp_allreduce, layout, false);
+                } else {
+                    reduce_lora_accumulator(
+                        ctx, ctx->grad_accum_a[offset + pair],
+                        scale * replicated_a2a_expert_scale,
+                        false, dp_allreduce);
+                    reduce_lora_accumulator(
+                        ctx, ctx->grad_accum_b[offset + pair],
+                        scale * replicated_a2a_expert_scale,
+                        false, dp_allreduce);
                 }
+                continue;
+            }
+            const auto layout = lora_tp_layout(ctx, layer, pair);
+            if (grouped_sync) {
+                grouped_plan.add(
+                    ctx, ctx->grad_accum_a[offset + pair], scale,
+                    at::Tensor(), at::Tensor(), 1.0,
+                    sharded_a2a, dp_allreduce, layout, true);
+                grouped_plan.add(
+                    ctx, ctx->grad_accum_b[offset + pair], scale,
+                    at::Tensor(), at::Tensor(), 1.0,
+                    sharded_a2a, dp_allreduce, layout, false);
+            } else {
+                reduce_lora_accumulator(
+                    ctx, ctx->grad_accum_a[offset + pair], scale,
+                    sharded_a2a, dp_allreduce);
+                reduce_lora_accumulator(
+                    ctx, ctx->grad_accum_b[offset + pair], scale,
+                    sharded_a2a, dp_allreduce);
+            }
+        }
+    }
+    if (grouped_sync) {
+        grouped_plan.execute(ctx);
+    } else {
+        for (int64_t layer = 0; layer < ctx->num_layers; ++layer) {
+            const int64_t offset = ctx->lora_layer_offset[layer];
+            const int64_t pairs = lora_pair_count(ctx->layer_configs[layer]);
+            for (int64_t pair = 0; pair < pairs; ++pair) {
+                const auto layout = lora_tp_layout(ctx, layer, pair);
+                tp_sum_replicated_lora_accumulator(
+                    ctx, ctx->grad_accum_a[offset + pair], layout, true);
+                tp_sum_replicated_lora_accumulator(
+                    ctx, ctx->grad_accum_b[offset + pair], layout, false);
             }
         }
+    }
+    cp_sum_fixed_lora_accumulators(ctx);
+    return true;
+}
 
-        // emptyCache for backward groups when seq>4096.
-        if (hidden_grad.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache();
+struct GradientClipEntry {
+    at::Tensor* accumulator;
+    int32_t group;
+    bool norm_owner;
+};
 
-        // Gradient for this group's input = gradient for next group's output
-        grad = grads[0];
+static bool gradient_norm_owner(
+    const TrainingContext* ctx, bool grouped_expert,
+    LoraTpLayout layout, bool is_a
+) {
+    // DP and CP keep complete LoRA tensors replicated. Count one owner before
+    // reducing the scalar over those axes. EP only replicates dense tensors;
+    // routed experts are disjoint across EP ranks. TP ownership follows the
+    // projection-aware LoRA layout (one factor replicated, one sharded).
+    if (ctx->data_parallel && ctx->dp_world_size > 1 && ctx->dp_rank != 0)
+        return false;
+    if (ctx->cp_world_size > 1 && ctx->cp_rank != 0)
+        return false;
+    if (ctx->expert_parallel && ctx->ep_world_size > 1 &&
+            !grouped_expert && ctx->ep_rank != 0)
+        return false;
+    if (ctx->tp_world_size > 1 &&
+            lora_tp_parameter_replicated(layout, is_a) && ctx->tp_rank != 0)
+        return false;
+    return true;
+}
 
-        // Free saved input to release memory
-        ctx->group_inputs[g] = at::Tensor();
-    }
+static void reduce_clip_norm_squares(
+    TrainingContext* ctx, at::Tensor& norm_squares
+) {
+    auto stream = c10::cuda::getCurrentCUDAStream(
+        norm_squares.device().index()).stream();
+    auto reduce_axis = [&](ncclComm_t communicator, int world_size,
+                           const char* axis) {
+        if (!communicator || world_size <= 1) return;
+        const auto error = ncclAllReduce(
+            norm_squares.data_ptr<float>(), norm_squares.data_ptr<float>(),
+            norm_squares.numel(), ncclFloat, ncclSum, communicator, stream);
+        TORCH_CHECK(error == ncclSuccess,
+            "NCCL ", axis, " gradient norm all-reduce failed: ",
+            ncclGetErrorString(error));
+    };
+    reduce_axis(ctx->tp_comm, ctx->tp_world_size, "TP");
+    reduce_axis(ctx->cp_comm, ctx->cp_world_size, "CP");
+    reduce_axis(ctx->nccl_comm, ctx->ep_world_size, "EP");
+    reduce_axis(ctx->dp_comm, ctx->dp_world_size, "DP");
+    auto pp_control = ctx->pp_control_comm ?
+        ctx->pp_control_comm : ctx->pp_comm;
+    reduce_axis(pp_control, ctx->pp_world_size, "PP");
 }
 
-// ──────────────────────────────────────────────────────────────────────
-// Fused Cross-Entropy Loss with online softmax (FlashAttention-style).
-//
-// Instead of materializing [n_tokens, vocab] logits (8+ GB), we tile over
-// the vocabulary dimension:
-//   1. Forward: iterate vocab tiles, compute partial logits, accumulate
-//      global max + sum_exp via online softmax → loss
-//   2. Backward: iterate vocab tiles again, compute softmax = exp(logit-max)/sum_exp,
-//      subtract one_hot for target tokens → accumulate grad into hidden_normed
-//
-// Peak memory: [n_tokens, tile_size] instead of [n_tokens, vocab].
-// For N=100, seq=512: n_tokens=51200, vocab=248320, tile=8192
-//   Old: [51200, 248320] × 4 bytes = 49 GB per chunk
-//   New: [51200, 8192]  × 4 bytes = 1.6 GB per tile
-//
-// Returns: scalar loss tensor (with autograd graph for hidden_normed)
-// ──────────────────────────────────────────────────────────────────────
-static at::Tensor compute_loss_fused(
+static void clip_lora_gradient_entries(
     TrainingContext* ctx,
-    const at::Tensor& hidden,       // [batch, seq, hidden] (requires_grad)
-    const at::Tensor& input_ids,    // [batch, seq]
-    const at::Tensor& target_mask,  // [batch, seq]
-    int64_t vocab_size
+    const std::vector<GradientClipEntry>& entries,
+    int32_t group_count
 ) {
-    auto final_norm = *ctx->final_norm_ptr[0];
-    auto lm_head = *ctx->lm_head_ptr[0];  // [vocab, hidden]
-
-    // Compute hidden_normed in no-grad, then set requires_grad.
-    auto hidden_detached = hidden.detach();
-    at::Tensor hidden_normed;
-    {
-        at::AutoGradMode no_grad_mode(false);
-        hidden_normed = rms_norm(hidden_detached, final_norm, ctx->rms_eps);
+    if (!ctx || ctx->max_grad_norm <= 0.0) return;
+    TORCH_CHECK(std::isfinite(ctx->max_grad_norm) &&
+            ctx->max_grad_norm > 0.0 && ctx->max_grad_norm <=
+            static_cast<double>(std::numeric_limits<float>::max()),
+        "native LoRA max gradient norm must be finite and representable as FP32");
+    TORCH_CHECK(group_count > 0, "native LoRA gradient clipping requires groups");
+    at::Tensor reference;
+    for (const auto& entry : entries) {
+        if (entry.accumulator && entry.accumulator->defined()) {
+            reference = *entry.accumulator;
+            break;
+        }
+    }
+    if (!reference.defined()) {
+        for (const auto* weight : ctx->weight_ptrs) {
+            if (weight && weight->defined()) {
+                reference = *weight;
+                break;
+            }
+        }
+    }
+    TORCH_CHECK(reference.defined() && reference.is_cuda(),
+        "native LoRA gradient clipping requires a CUDA reference tensor");
+
+    std::vector<void*> all_ptrs;
+    std::vector<int32_t> all_sizes;
+    std::vector<int32_t> all_groups;
+    std::vector<void*> owner_ptrs;
+    std::vector<int32_t> owner_sizes;
+    std::vector<int32_t> owner_groups;
+    all_ptrs.reserve(entries.size());
+    all_sizes.reserve(entries.size());
+    all_groups.reserve(entries.size());
+    for (const auto& entry : entries) {
+        TORCH_CHECK(entry.accumulator && entry.accumulator->defined() &&
+                entry.accumulator->is_cuda() && entry.accumulator->is_contiguous() &&
+                entry.accumulator->scalar_type() == at::kFloat &&
+                entry.accumulator->numel() <= std::numeric_limits<int32_t>::max(),
+            "native LoRA gradient clipping received an invalid accumulator");
+        all_ptrs.push_back(entry.accumulator->data_ptr());
+        all_sizes.push_back(static_cast<int32_t>(entry.accumulator->numel()));
+        all_groups.push_back(entry.group);
+        if (entry.norm_owner) {
+            owner_ptrs.push_back(entry.accumulator->data_ptr());
+            owner_sizes.push_back(static_cast<int32_t>(entry.accumulator->numel()));
+            owner_groups.push_back(entry.group);
+        }
     }
-    hidden_normed.set_requires_grad(true);
-
-    int64_t seq_len = hidden_normed.size(1);
-    int64_t hidden_dim = hidden_normed.size(2);
-
-    auto shifted_hidden = hidden_normed.narrow(1, 0, seq_len - 1);
-    auto shifted_targets = input_ids.narrow(1, 1, seq_len - 1).reshape({-1});
-    auto shifted_mask = target_mask.narrow(1, 1, seq_len - 1).reshape({-1});
-
-    int64_t total_tokens = shifted_targets.size(0);
-    auto hidden_flat = shifted_hidden.reshape({-1, hidden_dim});
-
-    // Ensure hidden_flat is contiguous for narrow + matmul
-    hidden_flat = hidden_flat.contiguous();
-
-    auto mask_f = shifted_mask.to(at::kFloat);
 
-    // ── Tile configuration ──
-    // tile_size controls vocab granularity. Larger = fewer iterations but more memory.
-    // [total_tokens, tile] × 4 bytes (FP32). 8192 → ~1.6 GB for 51200 tokens.
-    int64_t tile_size = 8192;
-    const char* ts_env = getenv("QWEN36_CE_TILE");
-    if (ts_env) { tile_size = atol(ts_env); if (tile_size < 1) tile_size = 8192; }
-    int64_t num_tiles = (vocab_size + tile_size - 1) / tile_size;
+    auto& buffers = ctx->adam_dev_bufs;
+    buffers.ensure(static_cast<int>(std::max(all_ptrs.size(), owner_ptrs.size())), reference);
+    buffers.ensure_groups(group_count, reference);
+    auto norm_squares = buffers.clip_norm_squares.narrow(0, 0, group_count);
+    norm_squares.zero_();
+    auto upload = [&](std::vector<void*>& pointers,
+                      std::vector<int32_t>& sizes,
+                      std::vector<int32_t>& groups) {
+        const int count = static_cast<int>(pointers.size());
+        if (count == 0) return count;
+        auto long_options = at::TensorOptions().dtype(at::kLong).device(at::kCPU);
+        auto int_options = at::TensorOptions().dtype(at::kInt).device(at::kCPU);
+        auto pointers_cpu = at::from_blob(
+            pointers.data(), {count}, long_options);
+        auto sizes_cpu = at::from_blob(sizes.data(), {count}, int_options);
+        auto groups_cpu = at::from_blob(groups.data(), {count}, int_options);
+        buffers.grads_buf.narrow(0, 0, count).copy_(pointers_cpu);
+        buffers.sizes_buf.narrow(0, 0, count).copy_(sizes_cpu);
+        buffers.groups_buf.narrow(0, 0, count).copy_(groups_cpu);
+        return count;
+    };
+    const int owner_count = upload(owner_ptrs, owner_sizes, owner_groups);
+    auto stream = c10::cuda::getCurrentCUDAStream(reference.device().index()).stream();
+    if (owner_count > 0) {
+        launch_fused_multi_tensor_l2_norm(
+            reinterpret_cast<void**>(buffers.grads_buf.data_ptr()),
+            buffers.sizes_buf.data_ptr<int>(), buffers.groups_buf.data_ptr<int>(),
+            owner_count, norm_squares.data_ptr<float>(), static_cast<void*>(stream));
+        const auto norm_launch_error = cudaGetLastError();
+        TORCH_CHECK(norm_launch_error == cudaSuccess,
+            "fused LoRA gradient norm launch failed: ",
+            cudaGetErrorString(norm_launch_error));
+    }
+    reduce_clip_norm_squares(ctx, norm_squares);
+
+    const int all_count = upload(all_ptrs, all_sizes, all_groups);
+    if (all_count > 0) {
+        launch_fused_multi_tensor_clip(
+            reinterpret_cast<void**>(buffers.grads_buf.data_ptr()),
+            buffers.sizes_buf.data_ptr<int>(), buffers.groups_buf.data_ptr<int>(),
+            all_count, norm_squares.data_ptr<float>(),
+            static_cast<float>(ctx->max_grad_norm), static_cast<void*>(stream));
+        const auto clip_launch_error = cudaGetLastError();
+        TORCH_CHECK(clip_launch_error == cudaSuccess,
+            "fused LoRA gradient clip launch failed: ",
+            cudaGetErrorString(clip_launch_error));
+    }
+}
 
-    // lm_head transpose: we need [hidden, vocab] for matmul
-    // lm_head is [vocab, hidden], so lm_head.t() is [hidden, vocab]
-    // We narrow on dim 0 of lm_head (vocab dim), then transpose.
-    // lm_head_w: [tile, hidden] → .t() → [hidden, tile]
-    // hidden_flat: [total_tokens, hidden] × [hidden, tile] → [total_tokens, tile]
+static void clip_fixed_lora_gradients(TrainingContext* ctx) {
+    if (!ctx || ctx->max_grad_norm <= 0.0) return;
+    std::vector<GradientClipEntry> entries;
+    for (int64_t layer = 0; layer < ctx->num_layers; ++layer) {
+        const auto table = lora_projection_table(ctx->layer_configs[layer]);
+        const int64_t offset = ctx->lora_layer_offset[layer];
+        for (int64_t pair = 0; pair < table.count; ++pair) {
+            const int64_t index = offset + pair;
+            if (!ctx->lora_active[index]) continue;
+            const auto layout = lora_tp_layout(ctx, layer, pair);
+            entries.push_back({&ctx->grad_accum_a[index], 0,
+                gradient_norm_owner(ctx, table.entries[pair].grouped_expert,
+                    layout, true)});
+            entries.push_back({&ctx->grad_accum_b[index], 0,
+                gradient_norm_owner(ctx, table.entries[pair].grouped_expert,
+                    layout, false)});
+        }
+    }
+    clip_lora_gradient_entries(ctx, entries, 1);
+}
 
-    // ── Forward pass: compute loss via online softmax ──
-    // For each token i: loss_i = log(sum_v exp(logit_iv)) - logit_i,target_i
-    // We compute in two phases:
-    //   Phase 1: find global max and sum_exp across all vocab tiles
-    //   Phase 2: compute loss = log(sum_exp) - target_logit (for masked tokens)
+static void clip_dynamic_lora_gradients(
+    TrainingContext* ctx, const std::vector<uint8_t>& adapter_has_global_tokens
+) {
+    if (!ctx || ctx->max_grad_norm <= 0.0) return;
+    TORCH_CHECK(adapter_has_global_tokens.size() == ctx->adapters.size(),
+        "dynamic LoRA gradient clipping activity mismatch");
+    std::vector<GradientClipEntry> entries;
+    for (size_t adapter_index = 0;
+         adapter_index < ctx->adapters.size(); ++adapter_index) {
+        if (!adapter_has_global_tokens[adapter_index]) continue;
+        auto& adapter = ctx->adapters[adapter_index];
+        for (const auto& [layer_idx, pairs] : adapter.params) {
+            const auto table = lora_projection_table(ctx->layer_configs[layer_idx]);
+            const auto accum_it = adapter.grad_accum.find(layer_idx);
+            TORCH_CHECK(accum_it != adapter.grad_accum.end() &&
+                    accum_it->second.size() == pairs.size(),
+                "dynamic LoRA gradient clipping layout mismatch");
+            for (size_t pair = 0; pair < pairs.size(); ++pair) {
+                if (!pairs[pair].first.requires_grad()) continue;
+                const auto layout = lora_tp_layout(
+                    ctx, layer_idx, static_cast<int64_t>(pair));
+                if (accum_it->second[pair][0].defined()) {
+                    entries.push_back({&accum_it->second[pair][0],
+                        static_cast<int32_t>(adapter_index),
+                        gradient_norm_owner(ctx, table.entries[pair].grouped_expert,
+                            layout, true)});
+                }
+                if (accum_it->second[pair][1].defined()) {
+                    entries.push_back({&accum_it->second[pair][1],
+                        static_cast<int32_t>(adapter_index),
+                        gradient_norm_owner(ctx, table.entries[pair].grouped_expert,
+                            layout, false)});
+                }
+            }
+        }
+    }
+    clip_lora_gradient_entries(
+        ctx, entries, static_cast<int32_t>(ctx->adapters.size()));
+}
 
-    // Running max and sum_exp per token
-    auto logit_max = at::full({total_tokens, 1}, -std::numeric_limits<float>::infinity(),
-        at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device()));
-    auto sum_exp = at::zeros({total_tokens, 1},
-        at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device()));
+static void elide_trivial_attention_mask(TrainingContext* ctx) {
+    ctx->attention_lengths = at::Tensor();
+    if (!ctx->attention_mask.defined() || ctx->attention_mask.numel() == 0) return;
+    // A padding-free batch can use SDPA's native causal fast path. This is one
+    // scalar synchronization per step, instead of materializing [B,S,S] in
+    // every full-attention layer.
+    if (at::all(ctx->attention_mask != 0).item<bool>()) {
+        ctx->attention_mask = at::Tensor();
+        return;
+    }
+    ctx->attention_lengths = (ctx->attention_mask != 0).to(at::kInt).sum(1)
+        .clamp(0, ctx->attention_mask.size(1)).to(at::kInt).contiguous();
+}
 
-    // Phase 1: accumulate max and sum_exp
-    for (int64_t t = 0; t < num_tiles; t++) {
-        int64_t v_start = t * tile_size;
-        int64_t v_end = std::min(v_start + tile_size, vocab_size);
-        int64_t v_n = v_end - v_start;
+// Linear attention carries recurrent state across the sequence. Until the
+// kernel accepts packed cu_seqlens, only full sequences and strict right
+// padding are safe; a 0 -> 1 transition would let padding state leak into a
+// later real token (left padding and internal holes).
+static void validate_linear_attention_mask(
+    TrainingContext* ctx, const at::Tensor& attention_mask
+) {
+    if (!ctx || !attention_mask.defined() || attention_mask.numel() == 0) return;
+    bool has_linear = false;
+    for (const auto& cfg : ctx->layer_configs) {
+        if (cfg.layer_type != 0) {
+            has_linear = true;
+            break;
+        }
+    }
+    if (!has_linear) {
+        for (const auto& cfg : ctx->mtp_layer_configs) {
+            if (cfg.layer_type != 0) {
+                has_linear = true;
+                break;
+            }
+        }
+    }
+    if (!has_linear) return;
+    TORCH_CHECK(attention_mask.dim() == 2,
+        "linear-attention mask must be [batch, seq]");
+    auto mask = attention_mask.to(at::kBool);
+    if (mask.size(1) <= 1) return;
+    auto leading = mask.narrow(1, 0, mask.size(1) - 1);
+    auto trailing = mask.narrow(1, 1, mask.size(1) - 1);
+    auto invalid_transition = leading.logical_not().logical_and(trailing);
+    TORCH_CHECK(!invalid_transition.any().item<bool>(),
+        "linear attention only supports full or strict right-padding masks; "
+        "left-padding/internal holes require packed cu_seqlens support");
+}
 
-        auto lm_head_tile = lm_head.narrow(0, v_start, v_n);  // [v_n, hidden]
-        // [total_tokens, hidden] × [hidden, v_n] → [total_tokens, v_n]
-        auto logits_tile = at::matmul(hidden_flat, lm_head_tile.t()).to(at::kFloat);
+// ── Multi-LoRA: concat all adapters' A/B, 2x GEMM ──
+// Pre-build cache of concatenated A/B per (layer, module) pair.
+// Called once at start of forward; reused across all layers.
 
-        // Online softmax update: new_max = max(old_max, tile_max)
-        auto tile_max = std::get<0>(at::max(logits_tile, /*dim=*/1, /*keepdim=*/true));
-        auto new_max = at::max(logit_max, tile_max);
+static void precompute_lora_cache(TrainingContext* ctx) {
+    if (ctx->lora_cache_valid) return;
+    ctx->lora_cache.clear();
 
-        // Adjust sum_exp: exp(old - new_max) * old_sum + exp(tile - new_max) * tile_sum
-        auto old_exp = at::exp(logit_max - new_max);
-        auto tile_exp = at::exp(logits_tile - new_max);
-        sum_exp = old_exp * sum_exp + tile_exp.sum(/*dim=*/1, /*keepdim=*/true);
-        logit_max = new_max;
+    // Phase 1: collect all (cache_key, a_concat, b_concat) tuples
+    struct LoraEntry {
+        int64_t key;
+        at::Tensor a_concat;  // [sum_ranks, in]
+        at::Tensor b_concat;  // [out, sum_ranks]
+    };
+    std::vector<LoraEntry> entries;
 
-        c10::cuda::CUDACachingAllocator::emptyCache();
+    for (int64_t layer_idx = 0; layer_idx < ctx->num_layers; layer_idx++) {
+        int64_t num_pairs = lora_pair_count(ctx->layer_configs[layer_idx]);
+        for (int64_t pair_idx = 0; pair_idx < num_pairs; pair_idx++) {
+            auto projection_table = lora_projection_table(ctx->layer_configs[layer_idx]);
+            // Routed experts use activation-level low-rank GEMMs. Materializing
+            // B@A for every local expert would erase LoRA's memory advantage.
+            if (projection_table.entries[pair_idx].grouped_expert) continue;
+            std::vector<at::Tensor> a_list, b_list;
+            const char* module_name = lora_pair_name(ctx->layer_configs[layer_idx], pair_idx);
+            for (auto& adapter : ctx->adapters) {
+                if (!adapter.target_modules.empty() &&
+                    adapter.target_modules.find(module_name) == adapter.target_modules.end())
+                    continue;
+                if (!adapter.all_target_layers &&
+                    adapter.target_layers.find(layer_idx) ==
+                        adapter.target_layers.end())
+                    continue;
+                auto it = adapter.params.find(layer_idx);
+                if (it == adapter.params.end()) continue;
+                if (pair_idx >= (int64_t)it->second.size()) continue;
+                auto& [a, b] = it->second[pair_idx];
+                if (!a.requires_grad() && !b.requires_grad()) continue;
+                double scaling = adapter.alpha / (double)adapter.rank;
+                b_list.push_back(b * scaling);
+                a_list.push_back(a);
+            }
+            if (a_list.empty()) {
+                if (!ctx->lora_a.empty() && layer_idx < (int64_t)ctx->lora_layer_offset.size()) {
+                    int64_t la_offset = ctx->lora_layer_offset[layer_idx];
+                    if (la_offset + pair_idx < (int64_t)ctx->lora_a.size()) {
+                        if (!ctx->lora_active.empty() &&
+                            !ctx->lora_active[la_offset + pair_idx]) {
+                            continue;
+                        }
+                        b_list.push_back(ctx->lora_b[la_offset + pair_idx] * ctx->lora_scaling);
+                        a_list.push_back(ctx->lora_a[la_offset + pair_idx]);
+                    }
+                }
+            }
+            if (!a_list.empty()) {
+                at::Tensor a_concat, b_concat;
+                if (a_list.size() == 1) {
+                    a_concat = a_list[0];
+                    b_concat = b_list[0];
+                } else {
+                    a_concat = at::cat(a_list, 0);  // [sum_ranks, in]
+                    b_concat = at::cat(b_list, 1);  // [out, sum_ranks]
+                }
+                entries.push_back({lora_cache_key(layer_idx, pair_idx), a_concat, b_concat});
+            }
+        }
     }
 
-    // Phase 2: compute target logit for each token
-    // We need logits[target] — gather the target position's logit.
-    // Do a second pass to find target logit.
-    auto target_logit = at::zeros({total_tokens, 1},
-        at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device()));
-
+    // Phase 2: group by (out_dim, in_dim, sum_ranks) and batch matmul
+    // delta = b_concat @ a_concat  → [out, in]
+    // Group entries with identical shapes to use at::bmm
+    struct ShapeGroup {
+        int64_t out_dim, in_dim, sum_ranks;
+        std::vector<size_t> indices;
+    };
+    std::vector<ShapeGroup> groups;
+    for (size_t i = 0; i < entries.size(); i++) {
+        auto& e = entries[i];
+        int64_t out_dim = e.b_concat.size(0);
+        int64_t sum_ranks = e.b_concat.size(1);
+        int64_t in_dim = e.a_concat.size(1);
+        bool found = false;
+        for (auto& g : groups) {
+            if (g.out_dim == out_dim && g.in_dim == in_dim && g.sum_ranks == sum_ranks) {
+                g.indices.push_back(i);
+                found = true;
+                break;
+            }
+        }
+        if (!found) {
+            groups.push_back({out_dim, in_dim, sum_ranks, {i}});
+        }
+    }
+
+    // Phase 3: for each group, stack and bmm in BF16 (2x faster tensor cores)
+    // LoRA params are FP32 for gradient stability; cast to BF16 for matmul only.
+    // Autograd handles the cast backward (grad → FP32 automatically).
+    auto bf16 = at::kBFloat16;
+    for (auto& g : groups) {
+        int n = (int)g.indices.size();
+        if (n == 1) {
+            auto& e = entries[g.indices[0]];
+            auto delta = at::matmul(e.b_concat, e.a_concat);
+            ctx->lora_cache[e.key] = delta;  // already BF16
+        } else {
+            std::vector<at::Tensor> b_stack_vec, a_stack_vec;
+            b_stack_vec.reserve(n);
+            a_stack_vec.reserve(n);
+            for (auto idx : g.indices) {
+                b_stack_vec.push_back(entries[idx].b_concat);
+                a_stack_vec.push_back(entries[idx].a_concat);
+            }
+            auto b_stack = at::stack(b_stack_vec, 0);  // [N, out, sum_ranks] BF16
+            auto a_stack = at::stack(a_stack_vec, 0);  // [N, sum_ranks, in] BF16
+            auto deltas = at::bmm(b_stack, a_stack);   // [N, out, in] BF16
+            for (int i = 0; i < n; i++) {
+                ctx->lora_cache[entries[g.indices[i]].key] = deltas[i];
+            }
+        }
+    }
+
+    ctx->lora_cache_valid = true;
+}
+
+// ── Batched Multi-LoRA: activation-level B@(A@x) ──
+
+/// Prepare stacked A/B tensors for the requested adapter layer range.
+/// Stores in ctx->lora_batch_cache. Recompute groups pass their own range so
+/// manual checkpoint backward does not rebuild projections for other layers.
+/// Replaces precompute_lora_cache when N > 1.
+static void prepare_lora_batch(
+    TrainingContext* ctx, int64_t start_layer = 0, int64_t end_layer = -1,
+    const std::vector<size_t>* selected_indices = nullptr
+) {
+    ctx->lora_batch_cache.clear();
+    ctx->lora_batch_n = 0;
+    TORCH_CHECK(start_layer >= 0 && start_layer <= ctx->num_layers,
+        "invalid LoRA batch start layer: ", start_layer);
+    if (end_layer < 0) end_layer = ctx->num_layers;
+    TORCH_CHECK(end_layer >= start_layer && end_layer <= ctx->num_layers,
+        "invalid LoRA batch end layer: ", end_layer);
+
+    std::vector<size_t> adapter_order;
+    if (selected_indices) {
+        adapter_order = *selected_indices;
+    } else {
+        adapter_order.reserve(ctx->adapters.size());
+        for (size_t i = 0; i < ctx->adapters.size(); ++i)
+            adapter_order.push_back(i);
+    }
+    TORCH_CHECK(!adapter_order.empty(), "LoRA batch requires at least one adapter");
+    for (const auto index : adapter_order)
+        TORCH_CHECK(index < ctx->adapters.size(), "LoRA batch adapter index is out of range");
+    const size_t batch_adapter_count = adapter_order.size();
+    std::vector<double> adapter_scalings;
+    adapter_scalings.reserve(batch_adapter_count);
+    for (const auto index : adapter_order) {
+        const auto& adapter = ctx->adapters[index];
+        adapter_scalings.push_back(adapter.alpha / (double)adapter.rank);
+    }
+    if (adapter_scalings != ctx->lora_batch_scaling_values) {
+        ctx->lora_batch_scaling = at::Tensor();
+        ctx->lora_batch_scaling_values = adapter_scalings;
+    }
+
+    for (int64_t layer_idx = start_layer;
+         layer_idx < end_layer; layer_idx++) {
+        int64_t num_pairs = lora_pair_count(ctx->layer_configs[layer_idx]);
+        for (int64_t pair_idx = 0; pair_idx < num_pairs; pair_idx++) {
+            std::vector<at::Tensor> a_list, b_list;
+            std::vector<int64_t> active_indices;
+            std::vector<double> scalings;
+            const char* module_name = lora_pair_name(ctx->layer_configs[layer_idx], pair_idx);
+            if (ctx->pad_heterogeneous_lora_batch) {
+                using LoraPair = std::pair<at::Tensor, at::Tensor>;
+                std::vector<LoraPair*> active_pairs(
+                    batch_adapter_count, nullptr);
+                at::Tensor template_a;
+                at::Tensor template_b;
+                int64_t a_rank_dim = -1;
+                int64_t b_rank_dim = -1;
+                int64_t padded_rank = 0;
+
+                for (size_t ordinal = 0;
+                     ordinal < batch_adapter_count; ++ordinal) {
+                    const size_t adapter_index = adapter_order[ordinal];
+                    auto& adapter = ctx->adapters[adapter_index];
+                    if (!adapter.target_modules.empty() &&
+                        adapter.target_modules.find(module_name) ==
+                            adapter.target_modules.end())
+                        continue;
+                    if (!adapter.all_target_layers &&
+                        adapter.target_layers.find(layer_idx) ==
+                            adapter.target_layers.end())
+                        continue;
+                    auto it = adapter.params.find(layer_idx);
+                    if (it == adapter.params.end() ||
+                        pair_idx >= static_cast<int64_t>(it->second.size()))
+                        continue;
+                    auto& pair = it->second[pair_idx];
+                    auto& [a, b] = pair;
+                    if (!a.requires_grad() && !b.requires_grad()) continue;
+                    TORCH_CHECK(a.requires_grad() && b.requires_grad(),
+                        "heterogeneous LoRA A/B activity mismatch for layer ",
+                        layer_idx, " projection ", module_name,
+                        " adapter ", adapter.id);
+                    TORCH_CHECK(
+                        (a.dim() == 2 && b.dim() == 2) ||
+                            (a.dim() == 3 && b.dim() == 3),
+                        "heterogeneous LoRA expects paired dense or routed-expert "
+                        "tensors for layer ", layer_idx, " projection ",
+                        module_name, " adapter ", adapter.id,
+                        ": A=", a.sizes(), " B=", b.sizes());
+                    const int64_t current_a_rank_dim = a.dim() == 2 ? 0 : 1;
+                    const int64_t current_b_rank_dim = b.dim() == 2 ? 1 : 2;
+                    TORCH_CHECK(a.size(current_a_rank_dim) ==
+                            b.size(current_b_rank_dim),
+                        "heterogeneous LoRA A/B rank mismatch for layer ",
+                        layer_idx, " projection ", module_name,
+                        " adapter ", adapter.id);
+                    if (!template_a.defined()) {
+                        template_a = a;
+                        template_b = b;
+                        a_rank_dim = current_a_rank_dim;
+                        b_rank_dim = current_b_rank_dim;
+                    } else {
+                        TORCH_CHECK(a.dim() == template_a.dim() &&
+                                b.dim() == template_b.dim() &&
+                                current_a_rank_dim == a_rank_dim &&
+                                current_b_rank_dim == b_rank_dim,
+                            "heterogeneous LoRA tensor rank differs within layer ",
+                            layer_idx, " projection ", module_name);
+                        for (int64_t dim = 0; dim < a.dim(); ++dim) {
+                            if (dim != a_rank_dim) {
+                                TORCH_CHECK(a.size(dim) == template_a.size(dim),
+                                    "heterogeneous LoRA A geometry differs within layer ",
+                                    layer_idx, " projection ", module_name);
+                            }
+                        }
+                        for (int64_t dim = 0; dim < b.dim(); ++dim) {
+                            if (dim != b_rank_dim) {
+                                TORCH_CHECK(b.size(dim) == template_b.size(dim),
+                                    "heterogeneous LoRA B geometry differs within layer ",
+                                    layer_idx, " projection ", module_name);
+                            }
+                        }
+                    }
+                    padded_rank = std::max(
+                        padded_rank, a.size(current_a_rank_dim));
+                    active_pairs[ordinal] = &pair;
+                }
+                if (!template_a.defined()) continue;
+
+                auto padded_a_sizes = template_a.sizes().vec();
+                auto padded_b_sizes = template_b.sizes().vec();
+                padded_a_sizes[a_rank_dim] = padded_rank;
+                padded_b_sizes[b_rank_dim] = padded_rank;
+                a_list.reserve(batch_adapter_count);
+                b_list.reserve(batch_adapter_count);
+                active_indices.reserve(batch_adapter_count);
+                for (size_t ordinal = 0;
+                     ordinal < batch_adapter_count; ++ordinal) {
+                    auto* pair = active_pairs[ordinal];
+                    if (!pair) {
+                        a_list.push_back(at::zeros(
+                            padded_a_sizes, template_a.options()));
+                        b_list.push_back(at::zeros(
+                            padded_b_sizes, template_b.options()));
+                    } else {
+                        auto& [a, b] = *pair;
+                        const int64_t rank = a.size(a_rank_dim);
+                        if (rank == padded_rank) {
+                            a_list.push_back(a);
+                            b_list.push_back(b);
+                        } else {
+                            auto a_padding_sizes = a.sizes().vec();
+                            auto b_padding_sizes = b.sizes().vec();
+                            a_padding_sizes[a_rank_dim] = padded_rank - rank;
+                            b_padding_sizes[b_rank_dim] = padded_rank - rank;
+                            a_list.push_back(at::cat(
+                                {a, at::zeros(a_padding_sizes, a.options())},
+                                a_rank_dim));
+                            b_list.push_back(at::cat(
+                                {b, at::zeros(b_padding_sizes, b.options())},
+                                b_rank_dim));
+                        }
+                    }
+                    active_indices.push_back(static_cast<int64_t>(ordinal));
+                }
+            } else {
+                for (size_t ordinal = 0;
+                     ordinal < batch_adapter_count; ++ordinal) {
+                    const size_t adapter_index = adapter_order[ordinal];
+                    auto& adapter = ctx->adapters[adapter_index];
+                    if (!adapter.target_modules.empty() &&
+                        adapter.target_modules.find(module_name) ==
+                            adapter.target_modules.end())
+                        continue;
+                    if (!adapter.all_target_layers &&
+                        adapter.target_layers.find(layer_idx) ==
+                            adapter.target_layers.end())
+                        continue;
+                    auto it = adapter.params.find(layer_idx);
+                    if (it == adapter.params.end()) continue;
+                    if (pair_idx >= (int64_t)it->second.size()) continue;
+                    auto& [a, b] = it->second[pair_idx];
+                    if (!a.requires_grad() && !b.requires_grad()) continue;
+                    b_list.push_back(b);
+                    a_list.push_back(a);
+                    scalings.push_back(
+                        adapter.alpha / (double)adapter.rank);
+                    active_indices.push_back(static_cast<int64_t>(ordinal));
+                }
+            }
+            if (a_list.empty()) continue;
+
+            int64_t n = (int64_t)a_list.size();
+            if (ctx->lora_batch_n == 0) ctx->lora_batch_n = n;
+
+            auto a_stack = at::stack(a_list, 0);  // [N, rank, in]
+            auto b_stack = at::stack(b_list, 0);  // [N, out, rank]
+            at::Tensor scaling;
+            bool all_adapters_active = active_indices.size() ==
+                batch_adapter_count;
+            if (all_adapters_active) {
+                for (size_t i = 0; i < active_indices.size(); ++i) {
+                    if (active_indices[i] != (int64_t)i) {
+                        all_adapters_active = false;
+                        break;
+                    }
+                }
+            }
+            if (all_adapters_active) {
+                if (!ctx->lora_batch_scaling.defined()) {
+                    auto scaling_cpu = at::from_blob(
+                        adapter_scalings.data(),
+                        {(int64_t)adapter_scalings.size(), 1, 1},
+                        at::TensorOptions().dtype(at::kDouble)).clone();
+                    ctx->lora_batch_scaling = scaling_cpu.to(
+                        a_stack.device()).to(at::kBFloat16);
+                    ++ctx->lora_batch_scaling_upload_count;
+                }
+                scaling = ctx->lora_batch_scaling;
+            } else {
+                // Heterogeneous target subsets can omit adapters for a
+                // projection; retain the compact fallback for that case.
+                auto scaling_cpu = at::from_blob(
+                    scalings.data(), {(int64_t)n, 1, 1},
+                    at::TensorOptions().dtype(at::kDouble)).clone();
+                scaling = scaling_cpu.to(a_stack.device()).to(at::kBFloat16);
+            }
+
+            ctx->lora_batch_cache[lora_cache_key(layer_idx, pair_idx)] = {
+                a_stack, b_stack, scaling,
+                lora_tp_layout(ctx, layer_idx, pair_idx)
+            };
+            ++ctx->lora_batch_projection_build_count;
+        }
+    }
+    ctx->lora_batch_valid = true;
+}
+
+/// Build activation-level entries for the fixed adapter in TP mode. The
+/// singleton adapter dimension is expanded lazily to the input batch by
+/// lora_activation_delta/dynamic_expert_lora_delta.
+static void prepare_fixed_lora_batch(TrainingContext* ctx) {
+    ctx->lora_batch_cache.clear();
+    ctx->lora_batch_n = 1;
+    ctx->lora_batch_scaling = at::Tensor();
+    for (int64_t layer_idx = 0; layer_idx < ctx->num_layers; ++layer_idx) {
+        const int64_t pair_count = lora_pair_count(ctx->layer_configs[layer_idx]);
+        const int64_t offset = ctx->lora_layer_offset[layer_idx];
+        for (int64_t pair_idx = 0; pair_idx < pair_count; ++pair_idx) {
+            if (lora_projection_table(ctx->layer_configs[layer_idx])
+                    .entries[pair_idx].grouped_expert)
+                continue;
+            const int64_t slot = offset + pair_idx;
+            if (!legacy_lora_slot_active(ctx, slot)) continue;
+            auto& a = ctx->lora_a[slot];
+            auto& b = ctx->lora_b[slot];
+            auto scaling = at::full(
+                {1, 1, 1}, ctx->lora_scaling,
+                at::TensorOptions().dtype(a.scalar_type()).device(a.device()));
+            ctx->lora_batch_cache[lora_cache_key(layer_idx, pair_idx)] = {
+                a.unsqueeze(0), b.unsqueeze(0), scaling,
+                lora_tp_layout(ctx, layer_idx, pair_idx)};
+            ++ctx->lora_batch_projection_build_count;
+        }
+    }
+    ctx->lora_batch_valid = true;
+}
+
+/// Compute B@(A@x) * scaling — never materializes B@A.
+/// x: [N, seq, in], A: [N, rank, in], B: [N, out, rank], scaling: [N, 1, 1]
+/// returns: [N, seq, out]
+static at::Tensor lora_activation_delta(
+    TrainingContext* ctx,
+    const at::Tensor& x,          // [N, seq, in]
+    const at::Tensor& A,          // [N, rank, in]
+    const at::Tensor& B,          // [N, out, rank]
+    const at::Tensor& scaling,    // [N, 1, 1]
+    LoraTpLayout layout
+) {
+    // Cast to compute dtype (BF16)
+    auto kind = x.scalar_type();
+    auto A_c = A.to(kind);
+    auto B_c = B.to(kind);
+    auto s_c = scaling.to(kind);
+    if (A_c.size(0) == 1 && x.size(0) > 1) {
+        A_c = A_c.expand({x.size(0), A_c.size(1), A_c.size(2)});
+        B_c = B_c.expand({x.size(0), B_c.size(1), B_c.size(2)});
+        s_c = s_c.expand({x.size(0), 1, 1});
+    }
+    // Latent-rank TP sums local deltas in forward. Its replicated input must
+    // likewise sum the rank-local input-gradient contributions in backward;
+    // otherwise a later sharded LoRA branch feeds only a partial dgrad into
+    // preceding replicated layers.
+    auto lora_input = layout == LoraTpLayout::LatentRank
+        ? tp_copy_lora_input(ctx, x) : x;
+    // Ax = A @ x^T  → [N, rank, seq]
+    auto Ax = at::bmm(A_c, lora_input.transpose(-2, -1));
+    // delta = B @ Ax → [N, out, seq] → transpose → [N, seq, out]
+    auto delta = at::bmm(B_c, Ax).transpose(-2, -1);
+    auto scaled = delta * s_c;
+    return layout == LoraTpLayout::LatentRank
+        ? tp_allreduce_lora_delta(ctx, scaled) : scaled;
+}
+
+// Fuse the shared A@x stage for dense Q/K/V LoRA projections.  The base QKV
+// weight may already use QWEN36_FUSED_QKV, but activation-level LoRA still
+// has three independent A and B batched matmuls unless this opt-in path is
+// enabled.  Heterogeneous ranks/layouts deliberately fall back to the
+// projection-local implementation above.
+static bool lora_activation_deltas_shared_qkv(
+    TrainingContext* ctx, const at::Tensor& x,
+    const std::array<const LoraBatchEntry*, 3>& entries,
+    std::array<at::Tensor, 3>& outputs
+) {
+    if (!env_enabled("QWEN36_FUSED_LORA_QKV_A")) return false;
+    if (!x.defined() || x.dim() != 3) return false;
+    const auto* first = entries[0];
+    if (!first) return false;
+    const auto layout = first->layout;
+    const int64_t batch = x.size(0);
+    const int64_t input_size = x.size(2);
+    const int64_t rank = first->a_stack.dim() == 3
+        ? first->a_stack.size(1) : -1;
+    if (rank <= 0 || input_size <= 0) return false;
+
+    std::array<at::Tensor, 3> a_cast;
+    std::array<at::Tensor, 3> b_cast;
+    std::array<at::Tensor, 3> scaling_cast;
+    for (size_t index = 0; index < entries.size(); ++index) {
+        const auto* entry = entries[index];
+        if (!entry || entry->layout != layout || entry->a_stack.dim() != 3 ||
+                entry->b_stack.dim() != 3 || entry->scaling.dim() != 3 ||
+                entry->a_stack.device() != x.device() ||
+                entry->b_stack.device() != x.device() ||
+                entry->scaling.device() != x.device() ||
+                entry->a_stack.size(1) != rank ||
+                entry->a_stack.size(2) != input_size ||
+                entry->b_stack.size(2) != rank ||
+                (entry->a_stack.size(0) != 1 &&
+                    entry->a_stack.size(0) != batch) ||
+                (entry->b_stack.size(0) != 1 &&
+                    entry->b_stack.size(0) != batch) ||
+                (entry->scaling.size(0) != 1 &&
+                    entry->scaling.size(0) != batch) ||
+                entry->scaling.size(1) != 1 || entry->scaling.size(2) != 1) {
+            return false;
+        }
+        a_cast[index] = entry->a_stack.to(x.scalar_type());
+        b_cast[index] = entry->b_stack.to(x.scalar_type());
+        scaling_cast[index] = entry->scaling.to(x.scalar_type());
+        if (batch > 1) {
+            if (a_cast[index].size(0) == 1)
+                a_cast[index] = a_cast[index].expand({batch, rank, input_size});
+            if (b_cast[index].size(0) == 1)
+                b_cast[index] = b_cast[index].expand({batch,
+                    b_cast[index].size(1), rank});
+            if (scaling_cast[index].size(0) == 1)
+                scaling_cast[index] = scaling_cast[index].expand({batch, 1, 1});
+        }
+    }
+
+    // Latent-rank LoRA needs the same replicated input copy for all three
+    // projections.  Sharing this copy is important: otherwise the A-side
+    // fusion would still launch three identical input collectives.
+    auto lora_input = layout == LoraTpLayout::LatentRank
+        ? tp_copy_lora_input(ctx, x) : x;
+    auto a_cat = at::cat({a_cast[0], a_cast[1], a_cast[2]}, 1);
+    auto ax = at::bmm(a_cat, lora_input.transpose(-2, -1));
+    int64_t rank_offset = 0;
+    for (size_t index = 0; index < entries.size(); ++index) {
+        auto ax_part = ax.narrow(1, rank_offset, rank);
+        auto delta = at::bmm(b_cast[index], ax_part).transpose(-2, -1);
+        auto scaled = delta * scaling_cast[index];
+        outputs[index] = layout == LoraTpLayout::LatentRank
+            ? tp_allreduce_lora_delta(ctx, scaled) : scaled;
+        rank_offset += rank;
+    }
+    return true;
+}
+
+static const LoraBatchEntry* lora_batch_entry(
+    TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx
+) {
+    if (!ctx || !ctx->lora_batch_valid || pair_idx < 0) return nullptr;
+    auto it = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, pair_idx));
+    return it == ctx->lora_batch_cache.end() ? nullptr : &it->second;
+}
+
+static at::Tensor dense_mlp_forward_batched(
+    TrainingContext* ctx, int64_t layer_idx, const at::Tensor& hidden,
+    const at::Tensor& gate_proj, const at::Tensor& up_proj,
+    const at::Tensor& down_proj, at::ScalarType compute_type,
+    const LayerConfig* config_override
+) {
+    const auto& cfg = config_override
+        ? *config_override : ctx->layer_configs.at(layer_idx);
+    const int64_t gate_pair = lora_pair_index(cfg, "gate_proj");
+    const int64_t up_pair = lora_pair_index(cfg, "up_proj");
+    const int64_t down_pair = lora_pair_index(cfg, "down_proj");
+
+    auto mlp_input = tp_copy_base_mlp_input(ctx, hidden);
+    at::Tensor gate_out;
+    at::Tensor up_out;
+    const auto fused_fc1 = fused_mlp_fc1_weight(ctx, layer_idx);
+    if (fused_fc1.defined()) {
+        TORCH_CHECK(fused_fc1.dim() == 2 &&
+                fused_fc1.size(0) == gate_proj.size(0) + up_proj.size(0) &&
+                fused_fc1.size(1) == mlp_input.size(2),
+            "fused dense FC1 weight shape is incompatible with the local MLP geometry");
+        auto fc1 = at::matmul(mlp_input, fused_fc1.t());
+        gate_out = fc1.narrow(-1, 0, gate_proj.size(0));
+        up_out = fc1.narrow(-1, gate_proj.size(0), up_proj.size(0));
+    } else {
+        gate_out = at::matmul(mlp_input, gate_proj.t());
+        up_out = at::matmul(mlp_input, up_proj.t());
+    }
+    gate_out = add_batched_lora(
+        ctx, gate_out, mlp_input, lora_batch_entry(ctx, layer_idx, gate_pair));
+    up_out = add_batched_lora(
+        ctx, up_out, mlp_input, lora_batch_entry(ctx, layer_idx, up_pair));
+    auto activated = fused_swiglu_op(gate_out, up_out, 0.0);
+    auto result = at::matmul(activated, down_proj.t());
+    result = add_batched_lora(
+        ctx, result, activated, lora_batch_entry(ctx, layer_idx, down_pair));
+    // Base TP uses a row-parallel down projection.  Reduce its local hidden
+    // contribution before the residual add.
+    return tp_allreduce_base_mlp(ctx, result.to(compute_type));
+}
+
+__attribute__((noinline, visibility("default")))
+at::Tensor apply_multi_lora(
+    TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx,
+    const at::Tensor& base_weight
+) {
+    auto it = ctx->lora_cache.find(lora_cache_key(layer_idx, pair_idx));
+    if (it == ctx->lora_cache.end()) return base_weight;
+
+    // Cached delta_weight = b_concat @ a_concat (BF16, precomputed in batched bmm)
+    auto& delta = it->second;
+    return base_weight + delta;  // both BF16, no conversion needed
+}
+
+// Forward declarations for sub-layer checkpointing
+// Sub-layer checkpointing: split each layer into attn + mlp segments
+// Enabled by QWEN36_SUBCKPT=1 env var. Reduces peak memory by ~2x
+// at the cost of 2x extra recomputation per layer during backward.
+// ──────────────────────────────────────────────────────────────────────
+
+// Compute attention output from hidden
+at::Tensor compute_attn_only(
+    TrainingContext* ctx, const at::Tensor& hidden, int64_t layer_idx, at::ScalarType kind
+) {
+    const auto& cfg = ctx->layer_configs[layer_idx];
+    int64_t w_offset = 0;
+    for (int64_t j = 0; j < layer_idx; j++)
+        w_offset += weight_count_for_layer(ctx->layer_configs[j]);
+    auto attn_input = rms_norm(hidden, *ctx->weight_ptrs[w_offset + 0], cfg.rms_eps);
+
+    // Use batched path if lora_batch is active
+    if (ctx->lora_batch_valid) {
+        if (cfg.layer_type == 0) {
+            auto qp = *ctx->weight_ptrs[w_offset+2], qn = *ctx->weight_ptrs[w_offset+3];
+            auto kp = *ctx->weight_ptrs[w_offset+4], kn = *ctx->weight_ptrs[w_offset+5];
+            auto vp = *ctx->weight_ptrs[w_offset+6], op = *ctx->weight_ptrs[w_offset+7];
+            const at::Tensor fused_qkv =
+                fused_full_attention_qkv_weight(ctx, layer_idx);
+            return full_attention_batched(
+                ctx, attn_input, layer_idx, qp, qn, kp, kn, vp, op,
+                cfg.num_heads, cfg.num_kv_heads, cfg.head_dim,
+                cfg.partial_rotary_factor, cfg.rope_theta, cfg.rms_eps, kind,
+                ctx->attention_mask, fused_qkv);
+        } else {
+            auto qkv = *ctx->weight_ptrs[w_offset+2], z = *ctx->weight_ptrs[w_offset+3];
+            auto a = *ctx->weight_ptrs[w_offset+4], b = *ctx->weight_ptrs[w_offset+5];
+            auto al = *ctx->weight_ptrs[w_offset+6], db = *ctx->weight_ptrs[w_offset+7];
+            auto cw = *ctx->weight_ptrs[w_offset+8], nw = *ctx->weight_ptrs[w_offset+9];
+            auto op = *ctx->weight_ptrs[w_offset+10];
+            return linear_attention_batched(
+                ctx, attn_input, layer_idx, qkv, z, a, b, al, db, cw, nw, op,
+                cfg.num_k_heads, cfg.key_dim, cfg.num_v_heads, cfg.val_dim,
+                cfg.conv_kernel, cfg.rms_eps, kind, ctx->attention_lengths);
+        }
+    }
+
+    // Legacy path: weight-level LoRA
+    TORCH_CHECK(!ctx->base_tp_attention,
+        "base attention TP requires the activation-level LoRA path");
+    int64_t lora_count = lora_pair_count(cfg);
+    int64_t la_offset = ctx->lora_layer_offset[layer_idx];
+    bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size();
+    std::vector<at::Tensor*> la(lora_count, nullptr), lb(lora_count, nullptr);
+    if (has_lora) for (int64_t k = 0; k < lora_count; k++) { la[k] = &ctx->lora_a[la_offset + k]; lb[k] = &ctx->lora_b[la_offset + k]; }
+
+    if (cfg.layer_type == 0) {
+        auto qp = *ctx->weight_ptrs[w_offset+2], qn = *ctx->weight_ptrs[w_offset+3];
+        auto kp = *ctx->weight_ptrs[w_offset+4], kn = *ctx->weight_ptrs[w_offset+5];
+        auto vp = *ctx->weight_ptrs[w_offset+6], op = *ctx->weight_ptrs[w_offset+7];
+        if (has_lora) {
+            if (la[0]) qp = lora_delta(qp, *la[0], *lb[0], ctx->lora_scaling);
+            if (la[1]) kp = lora_delta(kp, *la[1], *lb[1], ctx->lora_scaling);
+            if (la[2]) vp = lora_delta(vp, *la[2], *lb[2], ctx->lora_scaling);
+            if (la[3]) op = lora_delta(op, *la[3], *lb[3], ctx->lora_scaling);
+        }
+        return full_attention(attn_input, qp, qn, kp, kn, vp, op,
+            cfg.num_heads, cfg.num_kv_heads, cfg.head_dim,
+            cfg.partial_rotary_factor, cfg.rope_theta, cfg.rms_eps, kind,
+            ctx->attention_mask);
+    } else {
+        auto qkv = *ctx->weight_ptrs[w_offset+2], z = *ctx->weight_ptrs[w_offset+3];
+        auto a = *ctx->weight_ptrs[w_offset+4], b = *ctx->weight_ptrs[w_offset+5];
+        auto al = *ctx->weight_ptrs[w_offset+6], db = *ctx->weight_ptrs[w_offset+7];
+        auto cw = *ctx->weight_ptrs[w_offset+8], nw = *ctx->weight_ptrs[w_offset+9];
+        auto op = *ctx->weight_ptrs[w_offset+10];
+        if (has_lora) {
+            if (la[0]) qkv = lora_delta(qkv, *la[0], *lb[0], ctx->lora_scaling);
+            if (la[1]) z = lora_delta(z, *la[1], *lb[1], ctx->lora_scaling);
+            if (la[2]) a = lora_delta(a, *la[2], *lb[2], ctx->lora_scaling);
+            if (la[3]) b = lora_delta(b, *la[3], *lb[3], ctx->lora_scaling);
+            if (la[4]) op = lora_delta(op, *la[4], *lb[4], ctx->lora_scaling);
+        }
+        return linear_attention(attn_input, qkv, z, a, b, al, db, cw, nw, op,
+            cfg.num_k_heads, cfg.key_dim, cfg.num_v_heads, cfg.val_dim,
+            cfg.conv_kernel, cfg.rms_eps, kind, ctx->attention_lengths);
+    }
+}
+
+// Compute MLP output from residual (hidden + attn_output)
+at::Tensor compute_mlp_only(
+    TrainingContext* ctx, const at::Tensor& residual, int64_t layer_idx, at::ScalarType kind
+) {
+    const auto& cfg = ctx->layer_configs[layer_idx];
+    int64_t w_offset = 0;
+    for (int64_t j = 0; j < layer_idx; j++)
+        w_offset += weight_count_for_layer(ctx->layer_configs[j]);
+    int64_t mlp_start = (cfg.layer_type == 0) ? 8 : 11;
+    auto post_attn = rms_norm(residual, *ctx->weight_ptrs[w_offset + 1], cfg.rms_eps);
+    if (cfg.num_experts > 0) {
+        const bool use_batched = ctx->lora_batch_valid;
+        const int64_t shared_gate_pair = lora_pair_index(cfg, "shared_gate_proj");
+        const int64_t shared_up_pair = lora_pair_index(cfg, "shared_up_proj");
+        const int64_t shared_down_pair = lora_pair_index(cfg, "shared_down_proj");
+        const int64_t expert_gate_up_pair = lora_pair_index(cfg, "experts_gate_up_proj");
+        const int64_t expert_down_pair = lora_pair_index(cfg, "experts_down_proj");
+        auto shared_gate = use_batched ? *ctx->weight_ptrs[w_offset+mlp_start+2]
+            : apply_multi_lora(ctx, layer_idx, shared_gate_pair,
+                *ctx->weight_ptrs[w_offset+mlp_start+2]);
+        auto shared_up = use_batched ? *ctx->weight_ptrs[w_offset+mlp_start+3]
+            : apply_multi_lora(ctx, layer_idx, shared_up_pair,
+                *ctx->weight_ptrs[w_offset+mlp_start+3]);
+        auto shared_down = use_batched ? *ctx->weight_ptrs[w_offset+mlp_start+4]
+            : apply_multi_lora(ctx, layer_idx, shared_down_pair,
+                *ctx->weight_ptrs[w_offset+mlp_start+4]);
+        auto expert_lora = routed_expert_lora(ctx, layer_idx, cfg);
+        return moe_forward(ctx, layer_idx,
+            cfg.nccl_comm, cfg.nccl_stream, post_attn,
+            *ctx->weight_ptrs[w_offset+mlp_start], *ctx->weight_ptrs[w_offset+mlp_start+1],
+            shared_gate, shared_up,
+            shared_down, *ctx->weight_ptrs[w_offset+mlp_start+5],
+            *ctx->weight_ptrs[w_offset+mlp_start+6],
+            expert_lora,
+            cfg.num_experts, cfg.top_k, cfg.moe_intermediate,
+            cfg.norm_topk_prob != 0, cfg.expert_start, cfg.expert_count,
+            kind, use_batched,
+            use_batched ? lora_batch_entry(ctx, layer_idx, shared_gate_pair) : nullptr,
+            use_batched ? lora_batch_entry(ctx, layer_idx, shared_up_pair) : nullptr,
+            use_batched ? lora_batch_entry(ctx, layer_idx, shared_down_pair) : nullptr,
+            use_batched ? lora_batch_entry(ctx, layer_idx, expert_gate_up_pair) : nullptr,
+            use_batched ? lora_batch_entry(ctx, layer_idx, expert_down_pair) : nullptr);
+    } else {
+        if (ctx->lora_batch_valid) {
+            return dense_mlp_forward_batched(
+                ctx, layer_idx, post_attn,
+                *ctx->weight_ptrs[w_offset+mlp_start],
+                *ctx->weight_ptrs[w_offset+mlp_start+1],
+                *ctx->weight_ptrs[w_offset+mlp_start+2], kind);
+        }
+        auto gate = apply_multi_lora(ctx, layer_idx,
+            lora_pair_index(cfg, "gate_proj"), *ctx->weight_ptrs[w_offset+mlp_start]);
+        auto up = apply_multi_lora(ctx, layer_idx,
+            lora_pair_index(cfg, "up_proj"), *ctx->weight_ptrs[w_offset+mlp_start+1]);
+        auto down = apply_multi_lora(ctx, layer_idx,
+            lora_pair_index(cfg, "down_proj"), *ctx->weight_ptrs[w_offset+mlp_start+2]);
+        auto mlp_input = tp_copy_base_mlp_input(ctx, post_attn);
+        return tp_allreduce_base_mlp(
+            ctx, dense_mlp_forward(mlp_input, gate, up, down, kind));
+    }
+}
+
+// Sub-layer checkpoint: wraps a function call with no-grad forward + recompute on backward
+struct SubLayerCkpt : public torch::autograd::Function<SubLayerCkpt> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx,
+        at::Tensor input, int64_t tc_val, int64_t layer_idx, bool is_attn
+    ) {
+        ctx->saved_data["tc"] = tc_val;
+        ctx->saved_data["layer"] = layer_idx;
+        ctx->saved_data["is_attn"] = is_attn;
+        bool offload = getenv("QWEN36_OFFLOAD_ACTIVATIONS");
+        if (offload) {
+            ctx->saved_data["input_cpu"] = input.detach().to(
+                at::TensorOptions().dtype(input.scalar_type()).device(at::kCPU).pinned_memory(true));
+            ctx->saved_data["device"] = input.device();
+        } else {
+            ctx->save_for_backward({input});
+        }
+        at::AutoGradMode guard(false);
+        auto* tc = reinterpret_cast<TrainingContext*>(tc_val);
+        if (is_attn) {
+            // Attn segment: returns hidden + attn_output (residual connection)
+            return input + compute_attn_only(tc, input, layer_idx, tc->compute_type);
+        } else {
+            // MLP segment: returns residual + mlp_out (full layer output)
+            return input + compute_mlp_only(tc, input, layer_idx, tc->compute_type);
+        }
+    }
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx, std::vector<at::Tensor> grad_output
+    ) {
+        at::Tensor input;
+        if (ctx->saved_data.count("input_cpu") > 0) {
+            input = ctx->saved_data["input_cpu"].toTensor().to(ctx->saved_data["device"].toDevice());
+        } else {
+            input = ctx->get_saved_variables()[0];
+        }
+        auto* tc = reinterpret_cast<TrainingContext*>(ctx->saved_data["tc"].toInt());
+        int64_t layer = ctx->saved_data["layer"].toInt();
+        bool is_attn = ctx->saved_data["is_attn"].toBool();
+        at::AutoGradMode guard(true);
+        input.set_requires_grad(true);
+
+        // Derived LoRA deltas belong to one recomputed segment's graph and are
+        // released by grad(..., retain_graph=false). Attention and MLP both
+        // need a fresh cache now that both segments can own adapters.
+        tc->lora_batch_valid = false;
+        tc->lora_cache_valid = false;
+        if (!tc->adapters.empty()) prepare_lora_batch(tc);
+        else if (tc->tp_world_size > 1 || tc->cp_world_size > 1)
+            prepare_fixed_lora_batch(tc);
+        else precompute_lora_cache(tc);
+        auto output = is_attn ? compute_attn_only(tc, input, layer, tc->compute_type)
+                              : compute_mlp_only(tc, input, layer, tc->compute_type);
+
+        // Collect all tensors to compute gradients for: input + LoRA params for this layer
+        // This way grad() accumulates gradients into LoRA params (leaf nodes) too.
+        auto projection_table = lora_projection_table(tc->layer_configs[layer]);
+        int64_t lora_count = projection_table.count;
+        int64_t la_offset = tc->lora_layer_offset[layer];
+        bool has_lora = (la_offset + lora_count) <= (int64_t)tc->lora_a.size();
+
+        std::vector<at::Tensor> grad_inputs = {input};
+        std::vector<std::pair<at::Tensor*, at::Tensor*>> active_params;
+        if (!tc->adapters.empty()) {
+            for (auto& adapter : tc->adapters) {
+                auto it = adapter.params.find(layer);
+                if (it == adapter.params.end()) continue;
+                for (int64_t k = 0; k < lora_count && k < (int64_t)it->second.size(); ++k) {
+                    auto segment = projection_table.entries[k].segment;
+                    if ((is_attn && segment != LoraSegment::Attention) ||
+                        (!is_attn && segment != LoraSegment::Mlp)) continue;
+                    auto& [a, b] = it->second[k];
+                    if (!a.requires_grad() && !b.requires_grad()) continue;
+                    active_params.push_back({&a, &b});
+                    grad_inputs.push_back(a);
+                    grad_inputs.push_back(b);
+                }
+            }
+        } else if (has_lora) {
+            for (int64_t k = 0; k < lora_count; k++) {
+                auto segment = projection_table.entries[k].segment;
+                if ((is_attn && segment != LoraSegment::Attention) ||
+                    (!is_attn && segment != LoraSegment::Mlp)) continue;
+                int64_t slot = la_offset + k;
+                if (!legacy_lora_slot_active(tc, slot)) continue;
+                active_params.push_back({&tc->lora_a[slot], &tc->lora_b[slot]});
+                grad_inputs.push_back(tc->lora_a[slot]);
+                grad_inputs.push_back(tc->lora_b[slot]);
+            }
+        }
+
+        auto grads = torch::autograd::grad(
+            {output}, grad_inputs, {grad_output[0]},
+            /*retain_graph=*/false, /*create_graph=*/false,
+            /*allow_unused=*/true
+        );
+
+        // Manually accumulate LoRA param gradients
+        if (!active_params.empty()) {
+            int64_t gi = 1;  // skip input grad (index 0)
+            for (auto& [param_a_ptr, param_b_ptr] : active_params) {
+                if (grads[gi].defined()) {
+                    auto& param_a = *param_a_ptr;
+                    if (param_a.grad().defined())
+                        param_a.grad().add_(grads[gi]);
+                    else
+                        param_a.mutable_grad() = grads[gi].clone();
+                }
+                gi++;
+                if (grads[gi].defined()) {
+                    auto& param_b = *param_b_ptr;
+                    if (param_b.grad().defined())
+                        param_b.grad().add_(grads[gi]);
+                    else
+                        param_b.mutable_grad() = grads[gi].clone();
+                }
+                gi++;
+            }
+        }
+
+        return {grads[0], at::Tensor(), at::Tensor(), at::Tensor()};
+    }
+};
+
+// Forward a single layer with sub-layer checkpointing
+// attn segment returns (hidden + attn_output) to avoid extra GPU tensor
+// mlp segment returns (residual + mlp_out) = full layer output
+at::Tensor forward_single_layer_subckpt(
+    TrainingContext* ctx, const at::Tensor& hidden, int64_t layer_idx
+) {
+    // attn segment: computes attn_output, returns hidden + attn_output
+    auto residual = SubLayerCkpt::apply(
+        hidden, (int64_t)(uintptr_t)ctx, layer_idx, true);
+    // mlp segment: computes mlp_out from residual, returns residual + mlp_out
+    auto result = SubLayerCkpt::apply(
+        residual, (int64_t)(uintptr_t)ctx, layer_idx, false);
+    return result;
+}
+
+// ── Batched attention variants (activation-level LoRA) ──
+// These wrap the base attention functions, applying LoRA delta as
+// B@(A@x) * scaling on the projected outputs instead of modifying weights.
+
+static at::Tensor full_attention_batched(
+    TrainingContext* ctx, const at::Tensor& hidden, int64_t layer_idx,
+    const at::Tensor& q_proj, const at::Tensor& q_norm,
+    const at::Tensor& k_proj, const at::Tensor& k_norm,
+    const at::Tensor& v_proj, const at::Tensor& o_proj,
+    int64_t num_heads, int64_t num_kv_heads, int64_t head_dim,
+    double partial_rotary_factor, double rope_theta,
+    double rms_eps, at::ScalarType kind,
+    const at::Tensor& attention_mask,
+    const at::Tensor& fused_qkv
+) {
+    // Compute Q/K/V with base weight, then add LoRA delta if present
+    auto projection_input = tp_copy_base_attention_input(ctx, hidden);
+    int64_t batch = projection_input.size(0);
+    int64_t seq = projection_input.size(1);
+    const bool use_context_parallel = context_parallel_enabled(ctx);
+    const bool use_cp_ring = use_context_parallel &&
+        env_enabled("QWEN36_CP_FULL_ATTENTION_RING");
+    if (use_context_parallel) {
+        const bool use_kv_gather =
+            env_enabled("QWEN36_CP_FULL_ATTENTION_KV_GATHER");
+        TORCH_CHECK(ctx->cp_comm && use_kv_gather != use_cp_ring &&
+                !ctx->base_tp_attention,
+            "full-attention context parallelism requires exactly one guarded "
+            "KV gather or ring path without tensor-parallel attention");
+        TORCH_CHECK(ctx->cp_world_size == 2 || use_cp_ring,
+            "full-attention CP_SIZE>2 requires ring attention");
+    }
+    if (ctx->base_tp_attention) {
+        TORCH_CHECK(num_heads % ctx->tp_world_size == 0 &&
+                num_kv_heads % ctx->tp_world_size == 0,
+            "full-attention heads must be divisible by TP_SIZE");
+        num_heads /= ctx->tp_world_size;
+        num_kv_heads /= ctx->tp_world_size;
+    }
+    int64_t qkv_dim = num_heads * head_dim;
+    const int64_t q_projection_dim = num_heads * head_dim * 2;
+    const int64_t k_projection_dim = num_kv_heads * head_dim;
+    const bool use_fused_qkv = fused_qkv.defined();
+    if (use_fused_qkv) {
+        TORCH_CHECK(fused_qkv.dim() == 2 &&
+                fused_qkv.size(0) == q_projection_dim +
+                    2 * k_projection_dim &&
+                fused_qkv.size(1) == projection_input.size(2),
+            "fused full-attention QKV weight shape is incompatible with "
+            "the local head geometry");
+    }
+
+    at::Tensor q;
+    at::Tensor k;
+    at::Tensor v;
+    if (use_fused_qkv) {
+        auto qkv = at::matmul(projection_input, fused_qkv.t());
+        q = qkv.narrow(-1, 0, q_projection_dim);
+        k = qkv.narrow(-1, q_projection_dim, k_projection_dim);
+        v = qkv.narrow(-1, q_projection_dim + k_projection_dim,
+            k_projection_dim);
+    } else {
+        q = at::matmul(projection_input, q_proj.t());
+        k = at::matmul(projection_input, k_proj.t());
+        v = at::matmul(projection_input, v_proj.t());
+    }
+
+    // Apply activation-level LoRA: q/k/v += B@(A@hidden) * scaling.  When
+    // all three projections share compatible tenant geometry, fuse their
+    // A-side batched matmul while keeping the output-specific B projections.
+    auto it_q = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 0));
+    auto it_k = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 1));
+    auto it_v = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 2));
+    const std::array<const LoraBatchEntry*, 3> qkv_entries = {
+        it_q == ctx->lora_batch_cache.end() ? nullptr : &it_q->second,
+        it_k == ctx->lora_batch_cache.end() ? nullptr : &it_k->second,
+        it_v == ctx->lora_batch_cache.end() ? nullptr : &it_v->second,
+    };
+    std::array<at::Tensor, 3> qkv_lora;
+    if (lora_activation_deltas_shared_qkv(
+            ctx, projection_input, qkv_entries, qkv_lora)) {
+        q = q + qkv_lora[0];
+        k = k + qkv_lora[1];
+        v = v + qkv_lora[2];
+    } else {
+        if (qkv_entries[0]) {
+            q = q + lora_activation_delta(ctx, projection_input,
+                qkv_entries[0]->a_stack, qkv_entries[0]->b_stack,
+                qkv_entries[0]->scaling, qkv_entries[0]->layout);
+        }
+        if (qkv_entries[1]) {
+            k = k + lora_activation_delta(ctx, projection_input,
+                qkv_entries[1]->a_stack, qkv_entries[1]->b_stack,
+                qkv_entries[1]->scaling, qkv_entries[1]->layout);
+        }
+        if (qkv_entries[2]) {
+            v = v + lora_activation_delta(ctx, projection_input,
+                qkv_entries[2]->a_stack, qkv_entries[2]->b_stack,
+                qkv_entries[2]->scaling, qkv_entries[2]->layout);
+        }
+    }
+
+    // Reshape Q: [batch, seq, num_heads, head_dim*2] → split into q and gate
+    q = q.view({batch, seq, num_heads, head_dim * 2});
+    auto qk = q.chunk(2, -1);
+    auto q_out = qk[0].transpose(1, 2);   // [batch, heads, seq, head_dim]
+    auto gate = qk[1].transpose(1, 2);
+
+    k = k.view({batch, seq, num_kv_heads, head_dim}).transpose(1, 2);
+    v = v.view({batch, seq, num_kv_heads, head_dim}).transpose(1, 2);
+
+    q_out = rms_norm(q_out, q_norm, rms_eps);
+    k = rms_norm(k, k_norm, rms_eps);
+
+    // Keep gate alive until after SDPA. The architecture applies it to the
+    // attention value before the output projection.
+
+    // RoPE
+    int64_t rotary_dim = (int64_t)(head_dim * partial_rotary_factor);
+    if (rotary_dim > 0) {
+        auto device = hidden.device();
+        const int64_t position_start = use_context_parallel
+            ? ctx->cp_rank * seq : 0;
+        auto pos = at::arange(
+            position_start, position_start + seq,
+            at::TensorOptions().dtype(at::kFloat).device(device)).unsqueeze(0);
+        auto exponents = at::arange(0, rotary_dim, 2, at::TensorOptions().dtype(at::kFloat).device(device)) / (double)rotary_dim;
+        auto inv_freq = (exponents * std::log(rope_theta)).exp().reciprocal();
+        auto freqs = pos.unsqueeze(-1) * inv_freq.unsqueeze(0);
+        auto emb = at::cat({freqs, freqs}, -1);
+        auto cos = emb.cos().unsqueeze(1).to(q_out.scalar_type());
+        auto sin = emb.sin().unsqueeze(1).to(q_out.scalar_type());
+        auto q_rot = q_out.narrow(-1, 0, rotary_dim);
+        auto k_rot = k.narrow(-1, 0, rotary_dim);
+        auto rotate_half_q = at::cat({-q_rot.narrow(-1, rotary_dim/2, rotary_dim/2), q_rot.narrow(-1, 0, rotary_dim/2)}, -1);
+        auto rotate_half_k = at::cat({-k_rot.narrow(-1, rotary_dim/2, rotary_dim/2), k_rot.narrow(-1, 0, rotary_dim/2)}, -1);
+        q_rot.mul_(cos).add_(rotate_half_q * sin);
+        k_rot.mul_(cos).add_(rotate_half_k * sin);
+        cos = at::Tensor(); sin = at::Tensor();
+    }
+
+    int64_t key_sequence = seq;
+    at::Tensor local_ring_kv;
+    if (use_cp_ring) {
+        local_ring_kv = at::cat({k, v}, -1).contiguous();
+        key_sequence = seq * ctx->cp_world_size;
+    } else if (use_context_parallel) {
+        auto local_k = k.transpose(1, 2).reshape({
+            batch, seq, k_projection_dim});
+        auto local_v = v.transpose(1, 2).reshape({
+            batch, seq, k_projection_dim});
+        auto global_kv = context_parallel_kv_gather(
+            ctx, at::cat({local_k, local_v}, -1).contiguous());
+        key_sequence = seq * ctx->cp_world_size;
+        TORCH_CHECK(global_kv.sizes() == at::IntArrayRef({
+                batch, key_sequence, k_projection_dim * 2}),
+            "full-attention CP KV gather produced an invalid shape");
+        k = global_kv.narrow(-1, 0, k_projection_dim)
+            .reshape({batch, key_sequence, num_kv_heads, head_dim})
+            .transpose(1, 2).contiguous();
+        v = global_kv.narrow(-1, k_projection_dim, k_projection_dim)
+            .reshape({batch, key_sequence, num_kv_heads, head_dim})
+            .transpose(1, 2).contiguous();
+    }
+
+    // GQA: no K/V expansion needed (PT 2.5+ enable_gqa=true)
+    double scale = 1.0 / std::sqrt((double)head_dim);
+
+    // SDPA with GQA
+    at::Tensor attn_out;
+    if (use_cp_ring) {
+        if (attention_mask.defined() && attention_mask.numel() > 0) {
+            TORCH_CHECK(attention_mask.dim() == 2 &&
+                    attention_mask.size(0) == batch &&
+                    attention_mask.size(1) == key_sequence,
+                "full-attention CP ring mask must be [batch, global_seq]");
+        }
+        auto ring_mask = attention_mask.defined()
+            ? attention_mask.to(at::kBool).contiguous()
+            : at::empty({0}, at::TensorOptions().dtype(at::kBool)
+                .device(q_out.device()));
+        attn_out = Qwen36RingAttentionFunction::apply(
+            q_out.contiguous(), local_ring_kv, ring_mask,
+            (int64_t)ctx->cp_comm,
+            (int64_t)reinterpret_cast<uintptr_t>(ctx->cp_stream),
+            ctx->cp_rank, ctx->cp_world_size,
+            num_heads / num_kv_heads, ctx->cp_rank * seq);
+    } else if (use_context_parallel) {
+        const int64_t query_start = ctx->cp_rank * seq;
+        auto query_positions = at::arange(
+            query_start, query_start + seq,
+            at::TensorOptions().dtype(at::kLong).device(q_out.device()))
+            .unsqueeze(1);
+        auto key_positions = at::arange(
+            key_sequence,
+            at::TensorOptions().dtype(at::kLong).device(q_out.device()))
+            .unsqueeze(0);
+        auto combined = key_positions.gt(query_positions)
+            .unsqueeze(0).unsqueeze(0);
+        if (attention_mask.defined() && attention_mask.numel() > 0) {
+            auto key_padding = attention_mask.to(at::kBool);
+            TORCH_CHECK(key_padding.dim() == 2 &&
+                    key_padding.size(0) == batch &&
+                    key_padding.size(1) == key_sequence,
+                "full-attention CP padding mask must be [batch, global_seq]");
+            combined = combined.logical_or(
+                key_padding.logical_not().unsqueeze(1).unsqueeze(1));
+        }
+        auto additive_mask = at::zeros(
+            {batch, 1, seq, key_sequence},
+            at::TensorOptions().dtype(q_out.scalar_type())
+                .device(q_out.device()));
+        additive_mask = additive_mask.masked_fill(
+            combined, -std::numeric_limits<float>::infinity());
+        attn_out = at::scaled_dot_product_attention(
+            q_out, k, v, additive_mask, 0.0, false, c10::nullopt, true);
+    } else if (attention_mask.defined() && attention_mask.numel() > 0) {
+        auto kpm = attention_mask.to(at::kBool);
+        while (kpm.dim() > 2) kpm = kpm.squeeze(0);
+        if (kpm.size(0) == 1) {
+            kpm = kpm.unsqueeze(1).unsqueeze(1).expand({batch, 1, 1, seq});
+        } else {
+            kpm = kpm.unsqueeze(1).unsqueeze(1);
+        }
+        // Combined causal + padding mask
+        auto causal = at::triu(at::ones({seq, seq}, at::TensorOptions().dtype(at::kBool).device(q_out.device())), 1);
+        causal = causal.unsqueeze(0).unsqueeze(0);  // [1, 1, S, S]
+        auto pad_mask = kpm.logical_not();  // [B, 1, 1, S]
+        auto combined = causal.logical_or(pad_mask);
+        auto additive_mask = at::zeros({batch, 1, seq, seq}, at::TensorOptions().dtype(q_out.scalar_type()).device(q_out.device()));
+        additive_mask = additive_mask.masked_fill(combined, -std::numeric_limits<float>::infinity());
+        attn_out = at::scaled_dot_product_attention(q_out, k, v, additive_mask, 0.0, false, c10::nullopt, true);
+    } else {
+        attn_out = at::scaled_dot_product_attention(q_out, k, v, c10::nullopt, 0.0, true, c10::nullopt, true);
+    }
+    auto gated_attn = attn_out * at::sigmoid(gate).to(attn_out.scalar_type());
+    auto attn_flat = gated_attn.transpose(1, 2).reshape({batch, seq, qkv_dim});
+    auto result = attn_flat.matmul(o_proj.t());
+
+    // Apply LoRA delta on o_proj output
+    auto it_o = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 3));
+    if (it_o != ctx->lora_batch_cache.end()) {
+        result = result + lora_activation_delta(ctx, attn_flat,
+            it_o->second.a_stack, it_o->second.b_stack, it_o->second.scaling,
+            it_o->second.layout);
+    }
+    return tp_allreduce_base_attention(ctx, result);
+}
+
+static at::Tensor linear_attention_batched(
+    TrainingContext* ctx, const at::Tensor& hidden, int64_t layer_idx,
+    const at::Tensor& in_proj_qkv, const at::Tensor& in_proj_z,
+    const at::Tensor& in_proj_a, const at::Tensor& in_proj_b,
+    const at::Tensor& a_log, const at::Tensor& dt_bias,
+    const at::Tensor& conv1d_w, const at::Tensor& norm_w, const at::Tensor& out_proj,
+    int64_t num_k_heads, int64_t key_dim, int64_t num_v_heads, int64_t val_dim,
+    int64_t conv_kernel, double rms_eps, at::ScalarType compute_type,
+    const at::Tensor& attention_lengths
+) {
+    // Full reimplementation of linear_attention non-chunked path with
+    // activation-level LoRA delta on QKV, Z, and out_proj.
+    auto projection_input = tp_copy_base_attention_input(ctx, hidden);
+    int64_t batch = projection_input.size(0);
+    int64_t seq = projection_input.size(1);
+    const int64_t local_sequence = seq;
+    const bool use_context_parallel = context_parallel_enabled(ctx);
+    if (base_tp_attention_enabled(ctx)) {
+        TORCH_CHECK(num_k_heads > 0 && num_v_heads > 0 &&
+                num_v_heads % num_k_heads == 0 &&
+                num_k_heads % ctx->tp_world_size == 0 &&
+                num_v_heads % ctx->tp_world_size == 0,
+            "linear-attention heads must preserve value-head groups and be divisible by TP_SIZE");
+        num_k_heads /= ctx->tp_world_size;
+        num_v_heads /= ctx->tp_world_size;
+    }
+    const int64_t projected_num_v_heads = num_v_heads;
+    const int64_t projected_v_size = projected_num_v_heads * val_dim;
+    int64_t q_size = num_k_heads * key_dim;
+    int64_t v_size = num_v_heads * val_dim;
+    int64_t qkv_dim = q_size * 2 + v_size;
+
+    // QKV projection + LoRA delta
+    auto qkv = at::matmul(projection_input, in_proj_qkv.t());
+    auto it_qkv = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 0));
+    if (it_qkv != ctx->lora_batch_cache.end()) {
+        qkv = qkv + lora_activation_delta(ctx, projection_input, it_qkv->second.a_stack,
+            it_qkv->second.b_stack, it_qkv->second.scaling, it_qkv->second.layout);
+    }
+
+    at::Tensor effective_conv1d_w = conv1d_w;
+    at::Tensor a;
+    at::Tensor b;
+    at::Tensor z;
+    at::Tensor cp_controls;
+    const bool use_fused_cp_exchange = use_context_parallel &&
+        env_enabled("QWEN36_GDN_FUSED_CP_EXCHANGE");
+    auto project_gdn_controls = [&]() {
+        const bool use_fused_ab = layer_idx >= 0 &&
+            layer_idx < static_cast<int64_t>(
+                ctx->fused_gdn_ab_weights.size()) &&
+            ctx->fused_gdn_ab_weights[layer_idx].defined();
+        if (use_fused_ab) {
+            auto ab = at::matmul(
+                projection_input, ctx->fused_gdn_ab_weights[layer_idx].t());
+            TORCH_CHECK(ab.size(-1) == projected_num_v_heads * 2,
+                "fused GDN A/B projection output shape mismatch");
+            a = ab.narrow(-1, 0, projected_num_v_heads);
+            b = ab.narrow(
+                -1, projected_num_v_heads, projected_num_v_heads);
+        } else {
+            a = at::matmul(projection_input, in_proj_a.t());
+            b = at::matmul(projection_input, in_proj_b.t());
+        }
+        auto it_a = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 2));
+        if (it_a != ctx->lora_batch_cache.end()) {
+            a = a + lora_activation_delta(ctx, projection_input,
+                it_a->second.a_stack, it_a->second.b_stack,
+                it_a->second.scaling, it_a->second.layout);
+        }
+        auto it_b = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 3));
+        if (it_b != ctx->lora_batch_cache.end()) {
+            b = b + lora_activation_delta(ctx, projection_input,
+                it_b->second.a_stack, it_b->second.b_stack,
+                it_b->second.scaling, it_b->second.layout);
+        }
+
+        z = at::matmul(projection_input, in_proj_z.t());
+        auto it_z = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 1));
+        if (it_z != ctx->lora_batch_cache.end()) {
+            z = z + lora_activation_delta(ctx, projection_input,
+                it_z->second.a_stack, it_z->second.b_stack,
+                it_z->second.scaling, it_z->second.layout);
+        }
+    };
+    if (use_fused_cp_exchange) project_gdn_controls();
+    if (use_context_parallel) {
+        TORCH_CHECK(!base_tp_attention_enabled(ctx) && ctx->cp_world_size == 2 &&
+                ctx->cp_comm && q_size % ctx->cp_world_size == 0 &&
+                v_size % ctx->cp_world_size == 0,
+            "GDN context parallelism requires CP2-divisible Q/K/V head bundles");
+        const int64_t local_q_size = q_size / ctx->cp_world_size;
+        const int64_t local_v_size = v_size / ctx->cp_world_size;
+        const int64_t local_v_heads = num_v_heads / ctx->cp_world_size;
+        const int64_t local_qkv_dim = local_q_size * 2 + local_v_size;
+        const int64_t local_control_dim = local_v_size + local_v_heads * 2;
+        std::vector<at::Tensor> peer_qkv;
+        std::vector<at::Tensor> peer_conv;
+        if (!use_fused_cp_exchange)
+            peer_qkv.reserve(ctx->cp_world_size);
+        peer_conv.reserve(ctx->cp_world_size);
+        for (int64_t peer = 0; peer < ctx->cp_world_size; ++peer) {
+            if (!use_fused_cp_exchange) {
+                peer_qkv.push_back(at::cat({
+                    qkv.narrow(-1, peer * local_q_size, local_q_size),
+                    qkv.narrow(-1, q_size + peer * local_q_size, local_q_size),
+                    qkv.narrow(-1, q_size * 2 + peer * local_v_size,
+                        local_v_size)}, -1));
+            }
+            peer_conv.push_back(at::cat({
+                conv1d_w.narrow(0, peer * local_q_size, local_q_size),
+                conv1d_w.narrow(0, q_size + peer * local_q_size,
+                    local_q_size),
+                conv1d_w.narrow(0, q_size * 2 + peer * local_v_size,
+                    local_v_size)}, 0));
+        }
+        if (use_fused_cp_exchange) {
+            std::vector<at::Tensor> payload_sections;
+            payload_sections.reserve(ctx->cp_world_size * 6);
+            for (int64_t peer = 0; peer < ctx->cp_world_size; ++peer) {
+                // One final cat preserves peer-major [Q,K,V,Z,A,B] layout;
+                // avoiding nested cats keeps the fused path from copying the
+                // complete payload more than once before NCCL.
+                payload_sections.push_back(
+                    qkv.narrow(-1, peer * local_q_size, local_q_size));
+                payload_sections.push_back(qkv.narrow(
+                    -1, q_size + peer * local_q_size, local_q_size));
+                payload_sections.push_back(qkv.narrow(
+                    -1, q_size * 2 + peer * local_v_size, local_v_size));
+                payload_sections.push_back(
+                    z.narrow(-1, peer * local_v_size, local_v_size));
+                payload_sections.push_back(
+                    a.narrow(-1, peer * local_v_heads, local_v_heads));
+                payload_sections.push_back(
+                    b.narrow(-1, peer * local_v_heads, local_v_heads));
+            }
+            auto exchanged = CpSequenceToHeadFunction::apply(
+                at::cat(payload_sections, -1).contiguous(),
+                (int64_t)ctx->cp_comm,
+                (int64_t)reinterpret_cast<uintptr_t>(ctx->cp_stream),
+                ctx->cp_world_size);
+            TORCH_CHECK(exchanged.sizes() == at::IntArrayRef({
+                    batch, seq * ctx->cp_world_size,
+                    local_qkv_dim + local_control_dim}),
+                "fused GDN context-parallel exchange produced an invalid shape");
+            qkv = exchanged.narrow(-1, 0, local_qkv_dim).contiguous();
+            cp_controls = exchanged.narrow(
+                -1, local_qkv_dim, local_control_dim).contiguous();
+        } else {
+            auto packed_qkv = at::cat(peer_qkv, -1).contiguous();
+            qkv = CpSequenceToHeadFunction::apply(
+                packed_qkv, (int64_t)ctx->cp_comm,
+                (int64_t)reinterpret_cast<uintptr_t>(ctx->cp_stream),
+                ctx->cp_world_size);
+        }
+        effective_conv1d_w = peer_conv[ctx->cp_rank].contiguous();
+        num_k_heads /= ctx->cp_world_size;
+        num_v_heads /= ctx->cp_world_size;
+        q_size = local_q_size;
+        v_size = local_v_size;
+        qkv_dim = q_size * 2 + v_size;
+        seq *= ctx->cp_world_size;
+        TORCH_CHECK(qkv.sizes() == at::IntArrayRef({batch, seq, qkv_dim}),
+            "GDN context-parallel QKV exchange produced an invalid shape");
+    }
+
+    // DIAG: dump after QKV projection
+    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
+        auto qkv_f = qkv.to(at::kFloat);
+        fprintf(stderr, "[diag-la] layer %ld qkv_proj: shape=[%ld,%ld,%ld] mean=%.6f std=%.6f [0,0,:5]=%.6f,%.6f,%.6f,%.6f,%.6f\n",
+                (long)layer_idx, (long)qkv_f.size(0), (long)qkv_f.size(1), (long)qkv_f.size(2),
+                qkv_f.mean().item<float>(), qkv_f.std().item<float>(),
+                qkv_f[0][0][0].item<float>(), qkv_f[0][0][1].item<float>(),
+                qkv_f[0][0][2].item<float>(), qkv_f[0][0][3].item<float>(),
+                qkv_f[0][0][4].item<float>());
+        // Dump weight layout info
+        auto w_f = in_proj_qkv.to(at::kFloat);
+        fprintf(stderr, "[diag-la] in_proj_qkv weight: shape=[%ld,%ld] mean=%.6f std=%.6f\n",
+                (long)w_f.size(0), (long)w_f.size(1), w_f.mean().item<float>(), w_f.std().item<float>());
+        auto conv_w_f = conv1d_w.to(at::kFloat);
+        fprintf(stderr, "[diag-la] conv1d weight: shape=[%ld,%ld,%ld] mean=%.6f std=%.6f\n",
+                (long)conv_w_f.size(0), (long)conv_w_f.size(1), (long)conv_w_f.size(2),
+                conv_w_f.mean().item<float>(), conv_w_f.std().item<float>());
+    }
+
+    auto qkv_t = qkv.transpose(1, 2);
+    int64_t pad = conv_kernel - 1;
+    auto padding = at::zeros({batch, qkv_dim, pad}, qkv.options());
+    auto padded = at::cat({padding, qkv_t}, 2);
+    auto conv_out = at::conv1d(padded, effective_conv1d_w, /*bias=*/{},
+        at::IntArrayRef({1}), at::IntArrayRef({0}), at::IntArrayRef({1}), qkv_dim);
+    conv_out = at::silu(conv_out.narrow(2, 0, seq));
+    auto qkv_conv = conv_out.transpose(1, 2);
+
+    // DIAG: dump after conv1d
+    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
+        auto qc_f = qkv_conv.to(at::kFloat);
+        fprintf(stderr, "[diag-la] layer %ld after_conv1d: mean=%.6f std=%.6f [0,0,:5]=%.6f,%.6f,%.6f,%.6f,%.6f\n",
+                (long)layer_idx, qc_f.mean().item<float>(), qc_f.std().item<float>(),
+                qc_f[0][0][0].item<float>(), qc_f[0][0][1].item<float>(),
+                qc_f[0][0][2].item<float>(), qc_f[0][0][3].item<float>(),
+                qc_f[0][0][4].item<float>());
+    }
+
+    // Flat QKV split (matches transformers Qwen3_5GatedDeltaNet.forward)
+    // in_proj_qkv outputs flat layout: [Q_all(2048) | K_all(2048) | V_all(4096)]
+    // NOT per-head interleaved. This matches the non-batched path (line ~517).
+    int64_t head_k_dim = key_dim;                        // 128 (already per-head)
+    int64_t head_v_dim = val_dim;                        // 128 (already per-head)
+    int64_t q_total = num_k_heads * head_k_dim;          // 2048
+    int64_t v_total = num_v_heads * head_v_dim;          // 4096
+    auto q = qkv_conv.narrow(-1, 0, q_total).reshape({batch, seq, num_k_heads, head_k_dim});
+    auto k = qkv_conv.narrow(-1, q_total, q_total).reshape({batch, seq, num_k_heads, head_k_dim});
+    auto v = qkv_conv.narrow(-1, q_total * 2, v_total).reshape({batch, seq, num_v_heads, head_v_dim});
+
+    // DIAG: dump Q/K/V after per-head split
+    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
+        auto q_f = q.to(at::kFloat);
+        auto k_f = k.to(at::kFloat);
+        auto v_f = v.to(at::kFloat);
+        fprintf(stderr, "[diag-la] after_split q: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
+                q_f.mean().item<float>(), q_f.std().item<float>(),
+                q_f[0][0][0][0].item<float>(), q_f[0][0][0][1].item<float>(), q_f[0][0][0][2].item<float>());
+        fprintf(stderr, "[diag-la] after_split k: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
+                k_f.mean().item<float>(), k_f.std().item<float>(),
+                k_f[0][0][0][0].item<float>(), k_f[0][0][0][1].item<float>(), k_f[0][0][0][2].item<float>());
+        fprintf(stderr, "[diag-la] after_split v: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
+                v_f.mean().item<float>(), v_f.std().item<float>(),
+                v_f[0][0][0][0].item<float>(), v_f[0][0][0][1].item<float>(), v_f[0][0][0][2].item<float>());
+    }
+
+    if (!use_fused_cp_exchange) project_gdn_controls();
+
+    at::Tensor effective_a_log = a_log;
+    at::Tensor effective_dt_bias = dt_bias;
+    if (use_context_parallel) {
+        at::Tensor controls;
+        if (use_fused_cp_exchange) {
+            controls = cp_controls;
+        } else {
+            std::vector<at::Tensor> peer_controls;
+            peer_controls.reserve(ctx->cp_world_size);
+            for (int64_t peer = 0; peer < ctx->cp_world_size; ++peer) {
+                peer_controls.push_back(at::cat({
+                    z.narrow(-1, peer * v_size, v_size),
+                    a.narrow(-1, peer * num_v_heads, num_v_heads),
+                    b.narrow(-1, peer * num_v_heads, num_v_heads)}, -1));
+            }
+            controls = CpSequenceToHeadFunction::apply(
+                at::cat(peer_controls, -1).contiguous(),
+                (int64_t)ctx->cp_comm,
+                (int64_t)reinterpret_cast<uintptr_t>(ctx->cp_stream),
+                ctx->cp_world_size);
+        }
+        TORCH_CHECK(controls.size(0) == batch && controls.size(1) == seq &&
+                controls.size(2) == v_size + num_v_heads * 2,
+            "GDN context-parallel control exchange produced an invalid shape");
+        z = controls.narrow(-1, 0, v_size);
+        a = controls.narrow(-1, v_size, num_v_heads);
+        b = controls.narrow(-1, v_size + num_v_heads, num_v_heads);
+        effective_a_log = a_log.narrow(
+            0, ctx->cp_rank * num_v_heads, num_v_heads);
+        effective_dt_bias = dt_bias.narrow(
+            0, ctx->cp_rank * num_v_heads, num_v_heads);
+    } else {
+        TORCH_CHECK(z.size(-1) == projected_v_size,
+            "GDN Z projection output shape mismatch");
+    }
+    z = z.reshape({batch, seq, num_v_heads, head_v_dim});
+
+    // g = -exp(A_log) * softplus(a + dt_bias)
+    auto a_log_f = effective_a_log.to(at::kFloat);
+    auto dt_bias_f = effective_dt_bias.to(at::kFloat);
+    auto a_f = a.to(at::kFloat);
+    auto g = a_log_f.unsqueeze(0).unsqueeze(0).exp().neg() * at::softplus(a_f + dt_bias_f.unsqueeze(0).unsqueeze(0));
+    auto beta = at::sigmoid(b);
+
+    // Expand Q/K to num_v_heads
+    int64_t n_rep = num_v_heads / num_k_heads;
+    q = q.repeat_interleave(n_rep, 2);
+    k = k.repeat_interleave(n_rep, 2);
+
+    // L2 normalize Q, K (per-head, matching HF)
+    auto q_f = q.to(at::kFloat);
+    auto k_f = k.to(at::kFloat);
+    q = q_f * (q_f.pow(2).sum(-1, true) + 1e-6).rsqrt();
+    k = k_f * (k_f.pow(2).sum(-1, true) + 1e-6).rsqrt();
+
+    // DIAG: dump after L2 norm
+    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
+        auto q_f = q.to(at::kFloat);
+        auto k_f = k.to(at::kFloat);
+        fprintf(stderr, "[diag-la] after_l2norm q: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
+                q_f.mean().item<float>(), q_f.std().item<float>(),
+                q_f[0][0][0][0].item<float>(), q_f[0][0][0][1].item<float>(), q_f[0][0][0][2].item<float>());
+        fprintf(stderr, "[diag-la] after_l2norm k: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
+                k_f.mean().item<float>(), k_f.std().item<float>(),
+                k_f[0][0][0][0].item<float>(), k_f[0][0][0][1].item<float>(), k_f[0][0][0][2].item<float>());
+        auto g_f = g.to(at::kFloat);
+        fprintf(stderr, "[diag-la] g: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
+                g_f.mean().item<float>(), g_f.std().item<float>(),
+                g_f[0][0][0].item<float>(), g_f[0][0][1].item<float>(), g_f[0][0][2].item<float>());
+        auto beta_f = beta.to(at::kFloat);
+        fprintf(stderr, "[diag-la] beta: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
+                beta_f.mean().item<float>(), beta_f.std().item<float>(),
+                beta_f[0][0][0].item<float>(), beta_f[0][0][1].item<float>(), beta_f[0][0][2].item<float>());
+    }
+
+    double scale = 1.0 / std::sqrt((double)head_k_dim);
+    q = q * scale;
+
+    auto q_t = q.transpose(1, 2).contiguous();
+    auto k_t = k.transpose(1, 2).contiguous();
+    auto v_t = v.to(at::kFloat).transpose(1, 2).contiguous();
+    auto g_t = g.transpose(1, 2).contiguous();
+    auto beta_t = beta.to(at::kFloat).transpose(1, 2).contiguous();
+
+    auto g_exp = g_t.exp();
+
+    // N-aware sub-batching: process N adapters in groups of 256 to limit
+    // state tensor size and improve SM occupancy.
+    // State is independent per adapter — safe to split by N dimension.
+    int64_t BH_total = batch * num_v_heads;
+    int64_t sub_batch = (BH_total > 8192) ? 256 : batch;  // 256 adapters or all if small
+    sub_batch = std::min(sub_batch, batch);
+
+    std::vector<at::Tensor> sub_outputs;
+    sub_outputs.reserve((batch + sub_batch - 1) / sub_batch);
+
+    for (int64_t sb = 0; sb < batch; sb += sub_batch) {
+        int64_t n = std::min(sub_batch, batch - sb);
+        int64_t BH = n * num_v_heads;
+
+        // Narrow on dim 0 (batch/adapter dimension), then reshape to [BH, seq, dim]
+        auto q_sub = q_t.narrow(0, sb, n);
+        auto k_sub = k_t.narrow(0, sb, n);
+        auto v_sub = v_t.narrow(0, sb, n);
+        auto g_sub = g_exp.narrow(0, sb, n);
+        auto beta_sub = beta_t.narrow(0, sb, n);
+        auto lengths_sub = attention_lengths.defined()
+            ? attention_lengths.narrow(0, sb, n).contiguous()
+            : at::Tensor();
+
+        auto q_contig = q_sub.reshape({BH, seq, head_k_dim}).contiguous().to(at::kFloat);
+        auto k_contig = k_sub.reshape({BH, seq, head_k_dim}).contiguous().to(at::kFloat);
+        auto v_contig = v_sub.reshape({BH, seq, head_v_dim}).contiguous().to(at::kFloat);
+        auto g_contig = g_sub.reshape({BH, seq}).contiguous().to(at::kFloat);
+        auto beta_contig = beta_sub.reshape({BH, seq}).contiguous().to(at::kFloat);
+        sub_outputs.push_back(GatedDeltaRuleFunction::apply(
+            q_contig, k_contig, v_contig, g_contig, beta_contig,
+            gdn_lengths_arg(lengths_sub, q_contig)));
+    }
+
+    auto outs = at::cat(sub_outputs, 0);
+
+    auto core_out = outs.reshape({batch, num_v_heads, seq, head_v_dim})
+                         .transpose(1, 2).to(compute_type);
+
+    // DIAG: dump after delta rule
+    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
+        auto co_f = core_out.to(at::kFloat);
+        fprintf(stderr, "[diag-la] after_delta_rule: mean=%.6f std=%.6f [0,0,0,:3]=%.6f,%.6f,%.6f\n",
+                co_f.mean().item<float>(), co_f.std().item<float>(),
+                co_f[0][0][0][0].item<float>(), co_f[0][0][0][1].item<float>(), co_f[0][0][0][2].item<float>());
+    }
+
+    auto core_flat = core_out.reshape({-1, head_v_dim});
+    auto z_flat = z.reshape({-1, head_v_dim});
+    auto variance = core_flat.to(at::kFloat).pow(2).mean(-1, true);
+    auto normed = (core_flat.to(at::kFloat) * (variance + rms_eps).rsqrt() * norm_w.to(at::kFloat)).to(core_flat.scalar_type());
+    auto gated = (normed * at::silu(z_flat.to(at::kFloat)).to(normed.scalar_type())).reshape({batch, seq, num_v_heads * head_v_dim});
+    if (use_context_parallel) {
+        gated = CpHeadToSequenceFunction::apply(
+            gated.contiguous(), (int64_t)ctx->cp_comm,
+            (int64_t)reinterpret_cast<uintptr_t>(ctx->cp_stream),
+            ctx->cp_world_size);
+        TORCH_CHECK(gated.sizes() == at::IntArrayRef({
+                batch, local_sequence, projected_v_size}),
+            "GDN context-parallel output exchange produced an invalid shape");
+    }
+    auto result = at::matmul(gated, out_proj.t());
+
+    // DIAG: dump after norm+gate+out_proj
+    if (getenv("QWEN36_DUMP_LAYERS") && layer_idx == 0) {
+        auto r_f = result.to(at::kFloat);
+        fprintf(stderr, "[diag-la] after_out_proj: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
+                r_f.mean().item<float>(), r_f.std().item<float>(),
+                r_f[0][0][0].item<float>(), r_f[0][0][1].item<float>(), r_f[0][0][2].item<float>());
+    }
+
+    // out_proj LoRA delta: result += B@(A@gated) * scaling
+    auto it_op = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 4));
+    if (it_op != ctx->lora_batch_cache.end()) {
+        result = result + lora_activation_delta(ctx, gated, it_op->second.a_stack,
+            it_op->second.b_stack, it_op->second.scaling, it_op->second.layout);
+    }
+
+    return tp_allreduce_base_attention(ctx, result);
+}
+// ──────────────────────────────────────────────────────────────────────
+
+static at::Tensor dense_mlp_forward(
+    const at::Tensor& hidden,
+    const at::Tensor& gate_proj, const at::Tensor& up_proj, const at::Tensor& down_proj,
+    at::ScalarType compute_type
+) {
+    int64_t batch = hidden.size(0), seq = hidden.size(1), hidden_dim = hidden.size(2);
+    auto flat = hidden.reshape({batch * seq, hidden_dim});
+    auto gate_out = at::matmul(flat, gate_proj.t());
+    auto up_out = at::matmul(flat, up_proj.t());
+    // Fused silu * up via Tilelang (falls back to ATen)
+    auto activated = fused_swiglu_op(gate_out, up_out, 0.0);  // Qwen3.6 dense MLP has no clamp
+    return at::matmul(activated, down_proj.t()).reshape({batch, seq, hidden_dim});
+}
+
+// Run the local stage layers starting from an already-materialized hidden
+// activation. Pipeline stages use this after receiving their predecessor's
+// activation; the full-model path keeps its existing embedding entry point.
+static at::Tensor forward_stage_layers(
+    TrainingContext* ctx,
+    const at::Tensor& initial_hidden
+) {
+    auto kind = ctx->compute_type;
+    at::Tensor hidden = initial_hidden;
+    for (int64_t i = 0; i < ctx->num_layers; ++i) {
+        int64_t w_offset = 0;
+        for (int64_t j = 0; j < i; ++j)
+            w_offset += weight_count_for_layer(ctx->layer_configs[j]);
+        const int64_t w_count = weight_count_for_layer(ctx->layer_configs[i]);
+        std::vector<at::Tensor*> layer_w(
+            ctx->weight_ptrs.begin() + w_offset,
+            ctx->weight_ptrs.begin() + w_offset + w_count);
+        hidden = forward_single_layer(
+            ctx, hidden, layer_w.data(), &ctx->layer_configs[i], i,
+            kind, ctx->attention_mask, ctx->attention_lengths,
+            ctx->lora_batch_valid);
+    }
+    return hidden;
+}
+
+// Forward pass (no checkpointing)
+static at::Tensor forward_full(
+    TrainingContext* ctx,
+    const at::Tensor& input_ids
+) {
+    if (!ctx->adapters.empty()) prepare_lora_batch(ctx);
+    else if (ctx->tp_world_size > 1 || ctx->cp_world_size > 1)
+        prepare_fixed_lora_batch(ctx);
+    else precompute_lora_cache(ctx);
+    auto kind = ctx->compute_type;
+    at::AutoGradMode guard(true);
+    at::Tensor hidden = vocabulary_embedding(ctx, input_ids);
+
+    // Debug: dump embedding output stats
+    if (getenv("QWEN36_DUMP_LAYERS")) {
+        auto h_f = hidden.to(at::kFloat);
+        fprintf(stderr, "[dump] embedding: mean=%.6f std=%.6f [0,:3]=%.6f,%.6f,%.6f\n",
+                h_f.mean().item<float>(), h_f.std().item<float>(),
+                h_f[0][0][0].item<float>(), h_f[0][0][1].item<float>(), h_f[0][0][2].item<float>());
+    }
+
+    for (int64_t i = 0; i < ctx->num_layers; i++) {
+        // Get weight pointers for this layer
+        int64_t w_offset = 0;
+        for (int64_t j = 0; j < i; j++)
+            w_offset += weight_count_for_layer(ctx->layer_configs[j]);
+        int64_t w_count = weight_count_for_layer(ctx->layer_configs[i]);
+        std::vector<at::Tensor*> layer_w(ctx->weight_ptrs.begin() + w_offset,
+                                         ctx->weight_ptrs.begin() + w_offset + w_count);
+
+        // Get LoRA pointers for this layer (nullptr if no LoRA for this layer)
+        int64_t lora_count = lora_pair_count(ctx->layer_configs[i]);
+        int64_t la_offset = ctx->lora_layer_offset[i];
+        // Check if this layer has LoRA params (la_offset < lora_a.size())
+        bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size();
+        std::vector<at::Tensor*> la_ptrs(lora_count, nullptr), lb_ptrs(lora_count, nullptr);
+        if (has_lora) {
+            for (int64_t k = 0; k < lora_count; k++) {
+                la_ptrs[k] = &ctx->lora_a[la_offset + k];
+                lb_ptrs[k] = &ctx->lora_b[la_offset + k];
+            }
+        }
+
+        hidden = forward_single_layer(ctx, hidden, layer_w.data(), &ctx->layer_configs[i], i,
+            kind, ctx->attention_mask, ctx->attention_lengths, ctx->lora_batch_valid);
+
+        // Debug: dump per-layer hidden state stats
+        if (getenv("QWEN36_DUMP_LAYERS")) {
+            auto h_f = hidden.to(at::kFloat);
+            fprintf(stderr, "[dump] layer %ld: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
+                    (long)i, h_f.mean().item<float>(), h_f.std().item<float>(),
+                    h_f[0][0][0].item<float>(), h_f[0][0][1].item<float>(), h_f[0][0][2].item<float>());
+        }
+
+        // No per-layer sync — let CUDA pipeline run asynchronously.
+        // emptyCache() here was the #1 cause of GPU underutilization (6% util).
+    }
+
+    return hidden;  // pre-norm hidden (for MTP)
+}
+
+// ──────────────────────────────────────────────────────────────────────
+// Gradient checkpointing: per-group recomputation
+// ──────────────────────────────────────────────────────────────────────
+
+// Run a group of layers forward (with grad enabled, for recomputation)
+static at::Tensor forward_layer_group(
+    TrainingContext* ctx,
+    const at::Tensor& input,
+    int64_t start_layer,
+    int64_t end_layer
+) {
+    auto kind = ctx->compute_type;
+    at::Tensor hidden = input;
+
+    // Normal path: full layer forward (sub-layer checkpointing is handled
+    // in forward_full_checkpoint, not here)
+    for (int64_t i = start_layer; i < end_layer; i++) {
+        int64_t w_offset = 0;
+        for (int64_t j = 0; j < i; j++)
+            w_offset += weight_count_for_layer(ctx->layer_configs[j]);
+        int64_t w_count = weight_count_for_layer(ctx->layer_configs[i]);
+        std::vector<at::Tensor*> layer_w(ctx->weight_ptrs.begin() + w_offset,
+                                         ctx->weight_ptrs.begin() + w_offset + w_count);
+
+        int64_t lora_count = lora_pair_count(ctx->layer_configs[i]);
+        int64_t la_offset = ctx->lora_layer_offset[i];
+        bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size();
+        std::vector<at::Tensor*> la_ptrs(lora_count, nullptr), lb_ptrs(lora_count, nullptr);
+        if (has_lora) {
+            for (int64_t k = 0; k < lora_count; k++) {
+                la_ptrs[k] = &ctx->lora_a[la_offset + k];
+                lb_ptrs[k] = &ctx->lora_b[la_offset + k];
+            }
+        }
+
+        hidden = forward_single_layer(ctx, hidden, layer_w.data(), &ctx->layer_configs[i], i,
+            kind, ctx->attention_mask, ctx->attention_lengths, ctx->lora_batch_valid);
+
+        // Debug: dump per-layer hidden state stats (also in checkpoint recompute path)
+        if (getenv("QWEN36_DUMP_LAYERS")) {
+            auto h_f = hidden.to(at::kFloat);
+            fprintf(stderr, "[dump] layer %ld: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
+                    (long)i, h_f.mean().item<float>(), h_f.std().item<float>(),
+                    h_f[0][0][0].item<float>(), h_f[0][0][1].item<float>(), h_f[0][0][2].item<float>());
+        }
+    }
+    return hidden;
+}
+
+// autograd::Function for checkpointing a group of layers.
+// Forward: run group WITHOUT grad (no intermediate activations stored).
+// Backward: recompute group WITH grad, then backprop through recomputed graph.
+struct GroupCheckpointFunction : public torch::autograd::Function<GroupCheckpointFunction> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx,
+        at::Tensor input,
+        int64_t tc_val,
+        int64_t start_layer,
+        int64_t end_layer
+    ) {
+        ctx->saved_data["tc"] = tc_val;
+        ctx->saved_data["start"] = start_layer;
+        ctx->saved_data["end"] = end_layer;
+
+        // Check if activation offload is enabled
+        bool offload = getenv("QWEN36_OFFLOAD_ACTIVATIONS");
+
+        if (offload) {
+            // Save input to CPU — frees GPU memory between groups
+            // We store the CPU copy in saved_data (not save_for_backward,
+            // because save_for_backward would keep it on GPU)
+            auto input_cpu = input.detach().to(at::TensorOptions().dtype(input.scalar_type()).device(at::kCPU).pinned_memory(true));
+            ctx->saved_data["input_cpu"] = input_cpu;
+            // Also store device for restoring later
+            ctx->saved_data["device"] = input.device();
+        } else {
+            ctx->save_for_backward({input});
+        }
+
+        // Run forward in NO-GRAD mode — intermediate activations are NOT stored.
+        at::AutoGradMode guard(false);
+        auto* tc = reinterpret_cast<TrainingContext*>(tc_val);
+        return forward_layer_group(tc, input, start_layer, end_layer);
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output
+    ) {
+        bool offload = ctx->saved_data.count("input_cpu") > 0;
+
+        at::Tensor input;
+        if (offload) {
+            // Restore input from CPU → GPU
+            auto input_cpu = ctx->saved_data["input_cpu"].toTensor();
+            auto device = ctx->saved_data["device"].toDevice();
+            input = input_cpu.to(device);
+        } else {
+            auto saved = ctx->get_saved_variables();
+            input = saved[0];
+        }
+
+        auto tc = reinterpret_cast<TrainingContext*>(ctx->saved_data["tc"].toInt());
+        int64_t start_layer = ctx->saved_data["start"].toInt();
+        int64_t end_layer = ctx->saved_data["end"].toInt();
+
+        // Recompute forward WITH grad enabled — builds autograd graph for LoRA params.
+        at::AutoGradMode guard(true);
+        input.set_requires_grad(true);
+        auto output = forward_layer_group(tc, input, start_layer, end_layer);
+
+        // Backprop through recomputed graph.
+        // retain_graph=false: each group's recomputed graph is independent.
+        // LoRA param gradients accumulate via autograd's accumulator (leaf nodes).
+        // The graph is freed immediately after backward — critical for memory.
+        torch::autograd::backward({output}, {grad_output[0]},
+            /*retain_graph=*/false, /*create_graph=*/false);
+        return {input.grad(), at::Tensor(), at::Tensor(), at::Tensor()};
+    }
+};
+
+// ──────────────────────────────────────────────────────────────────────
+// FusedLayerFunction: autograd::Function for single layer forward+backward.
+// Forward: run WITH grad (PyTorch saves intermediates in graph).
+// Backward: PyTorch autograd traverses graph — NO recompute needed.
+// This eliminates checkpoint recompute (the main bottleneck).
+// Controlled by QWEN36_FUSED_LAYER=1 env var.
+// ──────────────────────────────────────────────────────────────────────
+
+struct FusedLayerFunction : public torch::autograd::Function<FusedLayerFunction> {
+    static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx,
+        at::Tensor input,
+        int64_t tc_val,
+        int64_t layer_idx
+    ) {
+        ctx->saved_data["tc"] = tc_val;
+        ctx->saved_data["layer"] = layer_idx;
+        // No save_for_backward — PyTorch autograd graph handles it.
+        // Forward runs WITH grad — all intermediates saved in graph.
+        auto* tc = reinterpret_cast<TrainingContext*>(tc_val);
+        auto kind = tc->compute_type;
+
+        int64_t w_offset = 0;
+        for (int64_t j = 0; j < layer_idx; j++)
+            w_offset += weight_count_for_layer(tc->layer_configs[j]);
+        int64_t w_count = weight_count_for_layer(tc->layer_configs[layer_idx]);
+        std::vector<at::Tensor*> layer_w(tc->weight_ptrs.begin() + w_offset,
+                                         tc->weight_ptrs.begin() + w_offset + w_count);
+        int64_t lora_count = lora_pair_count(tc->layer_configs[layer_idx]);
+        int64_t la_offset = tc->lora_layer_offset[layer_idx];
+        bool has_lora = (la_offset + lora_count) <= (int64_t)tc->lora_a.size();
+        std::vector<at::Tensor*> la(lora_count, nullptr), lb(lora_count, nullptr);
+        if (has_lora) for (int64_t k = 0; k < lora_count; k++) {
+            la[k] = &tc->lora_a[la_offset + k]; lb[k] = &tc->lora_b[la_offset + k];
+        }
+        return forward_single_layer(tc, input, layer_w.data(), &tc->layer_configs[layer_idx],
+            layer_idx, kind, tc->attention_mask, tc->attention_lengths,
+            tc->lora_batch_valid);
+    }
+
+    static std::vector<at::Tensor> backward(
+        torch::autograd::AutogradContext* ctx,
+        std::vector<at::Tensor> grad_output
+    ) {
+        // PyTorch autograd handles backward through the graph built during forward.
+        // No recompute needed — just return grad_output as grad_input.
+        // The actual backward computation happens via PyTorch's autograd engine
+        // traversing the graph nodes (matmul backward, SDPA backward, etc.).
+        return {grad_output[0], at::Tensor(), at::Tensor()};
+    }
+};
+
+// Forward pass with fused layer (no checkpoint, no recompute).
+// Uses FusedLayerFunction per layer — PyTorch autograd handles backward.
+// QWEN36_FUSED_LAYER=1 enables this path.
+static at::Tensor forward_full_fused(
+    TrainingContext* ctx,
+    const at::Tensor& input_ids
+) {
+    if (!ctx->adapters.empty()) prepare_lora_batch(ctx);
+    else if (ctx->tp_world_size > 1 || ctx->cp_world_size > 1)
+        prepare_fixed_lora_batch(ctx);
+    else precompute_lora_cache(ctx);
+    at::Tensor hidden = vocabulary_embedding(ctx, input_ids);
+    hidden = hidden.detach().set_requires_grad(true);
+
+    for (int64_t i = 0; i < ctx->num_layers; i++) {
+        hidden = FusedLayerFunction::apply(
+            hidden,
+            (int64_t)(uintptr_t)ctx,
+            i
+        );
+    }
+
+    return hidden;
+}
+
+// Forward pass with gradient checkpointing — manual checkpoint (no autograd::Function)
+// Forward: no-grad, save group inputs (offloaded to CPU). Backward: manual recompute per group.
+// This avoids autograd engine retaining all group outputs simultaneously.
+static at::Tensor forward_full_checkpoint(
+    TrainingContext* ctx,
+    const at::Tensor& input_ids
+) {
+    // Use activation-level paths for dynamic adapters and LoRA TP; the
+    // legacy weight cache remains the single-rank fast path.
+    if (!ctx->adapters.empty()) {
+        prepare_lora_batch(ctx);
+    } else if (ctx->tp_world_size > 1 || ctx->cp_world_size > 1) {
+        prepare_fixed_lora_batch(ctx);
+    } else {
+        precompute_lora_cache(ctx);
+    }
+    at::Tensor hidden = vocabulary_embedding(ctx, input_ids);
+
+    if (getenv("QWEN36_DUMP_LAYERS")) {
+        auto h_f = hidden.to(at::kFloat);
+        fprintf(stderr, "[dump] ckpt embedding: mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
+                h_f.mean().item<float>(), h_f.std().item<float>(),
+                h_f[0][0][0].item<float>(), h_f[0][0][1].item<float>(), h_f[0][0][2].item<float>());
+    }
+
+    bool use_subckpt = env_enabled("QWEN36_SUBCKPT");
+
+    if (use_subckpt) {
+        at::AutoGradMode restore(true);
+        hidden = hidden.detach().set_requires_grad(true);
+        for (int64_t i = 0; i < ctx->num_layers; i++) {
+            hidden = forward_single_layer_subckpt(ctx, hidden, i);
+        }
+        return hidden;
+    }
+
+    // Group-level manual checkpointing with variable group size.
+    // Larger group_size = fewer recomputations in backward (faster) but more
+    // peak memory. Default gs=4 (from ctx->group_size), overridable via env.
+    at::AutoGradMode no_grad(false);
+    hidden = hidden.detach();
+
+    ctx->group_inputs.clear();
+    bool offload = getenv("QWEN36_OFFLOAD_ACTIVATIONS");
+
+    // Build group list using ctx->group_size (default 4).
+    // Env override: QWEN36_GROUP_SIZE=10 sets gs=10.
+    int64_t gs = ctx->group_size;
+    if (gs < 1) gs = 1;
+    const char* gs_env = getenv("QWEN36_GROUP_SIZE");
+    if (gs_env) { gs = atol(gs_env); if (gs < 1) gs = 1; }
+    if (env_enabled("QWEN36_TRAIN_TRACE")) {
+        fprintf(stderr, "[checkpoint] group_size=%ld (num_layers=%ld → %ld groups)\n",
+                (long)gs, (long)ctx->num_layers,
+                (long)((ctx->num_layers + gs - 1) / gs));
+    }
+
+    std::vector<std::pair<int64_t, int64_t>> groups;
+    for (int64_t i = 0; i < ctx->num_layers; i += gs) {
+        groups.push_back({i, std::min(i + gs, ctx->num_layers)});
+    }
+
+    // Save groups for backward
+    ctx->group_ranges = groups;
+
+    for (auto& [start, end] : groups) {
+        if (offload && start < groups.back().first) {
+            ctx->group_inputs.push_back(
+                hidden.to(at::TensorOptions().dtype(hidden.scalar_type()).device(at::kCPU).pinned_memory(true))
+            );
+        } else {
+            ctx->group_inputs.push_back(hidden.clone());
+        }
+
+        hidden = forward_layer_group(ctx, hidden, start, end);
+
+        if (getenv("QWEN36_DUMP_LAYERS")) {
+            auto h_f = hidden.to(at::kFloat);
+            fprintf(stderr, "[dump] ckpt group [%ld,%ld): mean=%.6f std=%.6f [0,0,:3]=%.6f,%.6f,%.6f\n",
+                    (long)start, (long)end, h_f.mean().item<float>(), h_f.std().item<float>(),
+                    h_f[0][0][0].item<float>(), h_f[0][0][1].item<float>(), h_f[0][0][2].item<float>());
+        }
+    }
+
+    // Return hidden on GPU with requires_grad for CE backward
+    at::AutoGradMode restore(true);
+    hidden = hidden.set_requires_grad(true);
+    return hidden;
+}
+
+struct PipelineLossResult {
+    double value;
+    at::Tensor hidden_grad;
+    double numerator = 0.0;
+    double denominator = 0.0;
+    at::Tensor sample_loss_numerators;
+    at::Tensor sample_token_counts;
+};
+
+static PipelineLossResult pipeline_compute_loss(
+    TrainingContext* ctx,
+    const at::Tensor& hidden,
+    const at::Tensor& input_ids,
+    const at::Tensor& target_mask
+);
+
+// Commit fixed-LoRA accumulators with one fused Adam launch. Ordinary and
+// pipeline execution share this boundary so their optimizer clocks and state
+// transitions remain identical.
+static bool apply_fixed_lora_optimizer(
+    TrainingContext* ctx,
+    const at::Tensor& target_mask,
+    double accumulated_token_weight
+) {
+    const bool has_global_tokens = synchronize_lora_gradients(
+        ctx, target_mask, accumulated_token_weight);
+    if (!has_global_tokens) {
+        clear_gradient_accumulators(ctx);
+        return false;
+    }
+
+    const bool local_gradients_finite =
+        fixed_lora_accumulators_are_finite(ctx);
+    const bool all_gradients_finite =
+        fixed_optimizer_collective_all_succeeded(
+            ctx, local_gradients_finite);
+    TORCH_CHECK(local_gradients_finite && all_gradients_finite,
+        "fixed LoRA optimizer rejected non-finite accumulated gradients");
+
+    clip_fixed_lora_gradients(ctx);
+
+    at::AutoGradMode guard(false);
+    ctx->lora_cache_valid = false;
+    ctx->lora_batch_valid = false;
+    const int64_t next_step = ctx->fixed_optimizer_step + 1;
+    const double step_f = static_cast<double>(next_step);
+    const double bias_correction1 = 1.0 - std::pow(ctx->beta1, step_f);
+    const double bias_correction2 = 1.0 - std::pow(ctx->beta2, step_f);
+    const double sqrt_bias_correction2 = std::sqrt(bias_correction2);
+    const float lr_scaled = static_cast<float>(
+        ctx->lr * sqrt_bias_correction2 / bias_correction1);
+    const float eps_scaled = static_cast<float>(ctx->eps * sqrt_bias_correction2);
+    const float one_minus_b1 = static_cast<float>(1.0 - ctx->beta1);
+    const float one_minus_b2 = static_cast<float>(1.0 - ctx->beta2);
+
+    std::vector<void*> h_params, h_grads;
+    std::vector<float*> h_m, h_v;
+    std::vector<int> h_sizes;
+    size_t adam_idx = 0;
+    for (size_t i = 0; i < ctx->lora_a.size(); ++i) {
+        auto& a = ctx->lora_a[i];
+        auto& a_grad = ctx->grad_accum_a[i];
+        if (ctx->lora_active[i] && a_grad.defined() &&
+            a.scalar_type() == at::kBFloat16) {
+            h_params.push_back(a.data_ptr());
+            h_grads.push_back(a_grad.data_ptr());
+            h_m.push_back(static_cast<float*>(ctx->adam_m[adam_idx].data_ptr()));
+            h_v.push_back(static_cast<float*>(ctx->adam_v[adam_idx].data_ptr()));
+            h_sizes.push_back(static_cast<int>(a.numel()));
+        }
+        ++adam_idx;
+        auto& b = ctx->lora_b[i];
+        auto& b_grad = ctx->grad_accum_b[i];
+        if (ctx->lora_active[i] && b_grad.defined() &&
+            b.scalar_type() == at::kBFloat16) {
+            h_params.push_back(b.data_ptr());
+            h_grads.push_back(b_grad.data_ptr());
+            h_m.push_back(static_cast<float*>(ctx->adam_m[adam_idx].data_ptr()));
+            h_v.push_back(static_cast<float*>(ctx->adam_v[adam_idx].data_ptr()));
+            h_sizes.push_back(static_cast<int>(b.numel()));
+        }
+        ++adam_idx;
+    }
+
+    const int n_params = static_cast<int>(h_params.size());
+    if (n_params > 0) {
+        TORCH_CHECK(!ctx->lora_a.empty(),
+            "fixed LoRA optimizer has gradients but no parameters");
+        ctx->adam_dev_bufs.ensure(n_params, ctx->lora_a[0]);
+        auto opts_cpu_long = at::TensorOptions().dtype(at::kLong).device(at::kCPU);
+        auto opts_cpu_int = at::TensorOptions().dtype(at::kInt).device(at::kCPU);
+        auto params_cpu = at::from_blob(h_params.data(), {n_params}, opts_cpu_long);
+        auto grads_cpu = at::from_blob(h_grads.data(), {n_params}, opts_cpu_long);
+        auto m_cpu = at::from_blob(h_m.data(), {n_params}, opts_cpu_long);
+        auto v_cpu = at::from_blob(h_v.data(), {n_params}, opts_cpu_long);
+        auto sizes_cpu = at::from_blob(h_sizes.data(), {n_params}, opts_cpu_int);
+        ctx->adam_dev_bufs.params_buf.narrow(0, 0, n_params).copy_(params_cpu);
+        ctx->adam_dev_bufs.grads_buf.narrow(0, 0, n_params).copy_(grads_cpu);
+        ctx->adam_dev_bufs.m_buf.narrow(0, 0, n_params).copy_(m_cpu);
+        ctx->adam_dev_bufs.v_buf.narrow(0, 0, n_params).copy_(v_cpu);
+        ctx->adam_dev_bufs.sizes_buf.narrow(0, 0, n_params).copy_(sizes_cpu);
+        auto stream = c10::cuda::getCurrentCUDAStream().stream();
+        launch_fused_adam_multi(
+            reinterpret_cast<void**>(ctx->adam_dev_bufs.params_buf.data_ptr()),
+            reinterpret_cast<void**>(ctx->adam_dev_bufs.grads_buf.data_ptr()),
+            reinterpret_cast<float**>(ctx->adam_dev_bufs.m_buf.data_ptr()),
+            reinterpret_cast<float**>(ctx->adam_dev_bufs.v_buf.data_ptr()),
+            reinterpret_cast<int*>(ctx->adam_dev_bufs.sizes_buf.data_ptr()),
+            n_params,
+            static_cast<float>(ctx->beta1), static_cast<float>(ctx->beta2),
+            lr_scaled, eps_scaled, one_minus_b1, one_minus_b2,
+            static_cast<void*>(stream));
+        auto launch_error = cudaGetLastError();
+        TORCH_CHECK(launch_error == cudaSuccess,
+            "fused FP32-gradient Adam launch failed: ",
+            cudaGetErrorString(launch_error));
+    }
+    // A pipeline stage may own no selected target slots. It still participates
+    // in the successful global optimizer commit and must advance its clock so
+    // PP checkpoint generations and later phase validation stay aligned.
+    ctx->fixed_optimizer_step = next_step;
+    clear_gradient_accumulators(ctx);
+    return true;
+}
+
+static void pipeline_send_tensor(
+    TrainingContext* ctx,
+    const at::Tensor& tensor,
+    int peer
+) {
+    TORCH_CHECK(ctx->pp_comm && peer >= 0 && peer < ctx->pp_world_size,
+        "pipeline send requires a valid PP communicator and peer");
+    auto stream = c10::cuda::getCurrentCUDAStream(
+        tensor.device().index()).stream();
+    TORCH_CHECK(ncclGroupStart() == ncclSuccess,
+        "pipeline forward ncclGroupStart failed");
+    auto send_error = ncclSend(
+        tensor.data_ptr(), tensor.numel(), qwen36_nccl_dtype(tensor.scalar_type()),
+        peer, ctx->pp_comm, stream);
+    auto end_error = ncclGroupEnd();
+    TORCH_CHECK(send_error == ncclSuccess && end_error == ncclSuccess,
+        "pipeline tensor send failed: ",
+        ncclGetErrorString(send_error == ncclSuccess ? end_error : send_error));
+}
+
+static at::Tensor pipeline_recv_tensor(
+    TrainingContext* ctx,
+    const at::Tensor& shape_source,
+    int peer,
+    int64_t hidden_size
+) {
+    TORCH_CHECK(ctx->pp_comm && peer >= 0 && peer < ctx->pp_world_size,
+        "pipeline receive requires a valid PP communicator and peer");
+    auto options = at::TensorOptions()
+        .device(shape_source.device())
+        .dtype(ctx->compute_type);
+    auto tensor = at::empty({shape_source.size(0), shape_source.size(1), hidden_size}, options);
+    auto stream = c10::cuda::getCurrentCUDAStream(
+        tensor.device().index()).stream();
+    TORCH_CHECK(ncclGroupStart() == ncclSuccess,
+        "pipeline backward ncclGroupStart failed");
+    auto recv_error = ncclRecv(
+        tensor.data_ptr(), tensor.numel(), qwen36_nccl_dtype(tensor.scalar_type()),
+        peer, ctx->pp_comm, stream);
+    auto end_error = ncclGroupEnd();
+    TORCH_CHECK(recv_error == ncclSuccess && end_error == ncclSuccess,
+        "pipeline tensor receive failed: ",
+        ncclGetErrorString(recv_error == ncclSuccess ? end_error : recv_error));
+    return tensor;
+}
+
+static at::Tensor pipeline_exchange_tensor(
+    TrainingContext* ctx,
+    const at::Tensor& send_tensor,
+    const at::Tensor& shape_source,
+    int peer,
+    int64_t hidden_size
+) {
+    TORCH_CHECK(ctx->pp_comm && peer >= 0 && peer < ctx->pp_world_size,
+        "pipeline exchange requires a valid PP communicator and peer");
+    TORCH_CHECK(send_tensor.is_cuda() && send_tensor.is_contiguous() &&
+            send_tensor.scalar_type() == ctx->compute_type &&
+            send_tensor.dim() == 3 &&
+            send_tensor.size(0) == shape_source.size(0) &&
+            send_tensor.size(1) == shape_source.size(1) &&
+            send_tensor.size(2) == hidden_size,
+        "pipeline exchange send tensor violates the fixed activation contract");
+    auto received = at::empty(
+        {shape_source.size(0), shape_source.size(1), hidden_size},
+        send_tensor.options());
+    auto stream = c10::cuda::getCurrentCUDAStream(
+        send_tensor.device().index()).stream();
+    TORCH_CHECK(ncclGroupStart() == ncclSuccess,
+        "pipeline exchange ncclGroupStart failed");
+    const auto send_error = ncclSend(
+        send_tensor.data_ptr(), send_tensor.numel(),
+        qwen36_nccl_dtype(send_tensor.scalar_type()),
+        peer, ctx->pp_comm, stream);
+    const auto recv_error = ncclRecv(
+        received.data_ptr(), received.numel(),
+        qwen36_nccl_dtype(received.scalar_type()),
+        peer, ctx->pp_comm, stream);
+    const auto end_error = ncclGroupEnd();
+    const auto error = send_error != ncclSuccess ? send_error
+        : (recv_error != ncclSuccess ? recv_error : end_error);
+    TORCH_CHECK(error == ncclSuccess,
+        "pipeline grouped tensor exchange failed: ", ncclGetErrorString(error));
+    return received;
+}
+
+static double pipeline_broadcast_loss(
+    TrainingContext* ctx,
+    double local_loss,
+    const at::TensorOptions& options
+) {
+    auto loss = at::full({1}, local_loss, options.dtype(at::kFloat));
+    auto stream = c10::cuda::getCurrentCUDAStream(
+        loss.device().index()).stream();
+    auto control_comm = ctx->pp_control_comm ?
+        ctx->pp_control_comm : ctx->pp_comm;
+    auto error = ncclBroadcast(
+        loss.data_ptr(), loss.data_ptr(), 1, ncclFloat32,
+        ctx->pp_world_size - 1, control_comm, stream);
+    TORCH_CHECK(error == ncclSuccess,
+        "pipeline loss broadcast failed: ", ncclGetErrorString(error));
+    return loss.to(at::kCPU).item<double>();
+}
+
+static void pipeline_window_write_result(
+    TrainingContext* ctx,
+    Qwen36PipelineResultV1* result,
+    int32_t status,
+    double loss
+) {
+    if (!result) return;
+    result->struct_size = sizeof(Qwen36PipelineResultV1);
+    result->version = 1;
+    result->status = status;
+    result->completed_fwd = ctx ? ctx->pp_window.next_forward : 0;
+    result->completed_bwd = ctx ? ctx->pp_window.next_backward : 0;
+    result->in_flight = ctx
+        ? static_cast<int64_t>(ctx->pp_window.slots.size()) : 0;
+    // v1 exposes one scalar clock for the fixed-LoRA path. Dynamic windows
+    // have one independent clock per tenant, so mark this field unavailable.
+    result->optimizer_step = ctx && !ctx->pp_window.dynamic_lora
+        ? ctx->fixed_optimizer_step : -1;
+    result->loss = loss;
+}
+
+static void pipeline_window_validate_result(
+    const Qwen36PipelineResultV1* result
+) {
+    if (!result) return;
+    TORCH_CHECK(result->struct_size >= sizeof(Qwen36PipelineResultV1) &&
+            result->version == 1,
+        "unsupported pipeline result ABI or undersized result buffer");
+}
+
+static bool pipeline_window_validate_forward_contract(
+    TrainingContext* ctx,
+    const Qwen36PipelineTickV1& tick
+) {
+    auto& window = ctx->pp_window;
+    const auto* input_ids = tick.input_ids
+        ? reinterpret_cast<const at::Tensor*>(tick.input_ids) : nullptr;
+    const auto* target_mask = tick.target_mask
+        ? reinterpret_cast<const at::Tensor*>(tick.target_mask) : nullptr;
+    const auto* attention_mask = tick.attention_mask
+        ? reinterpret_cast<const at::Tensor*>(tick.attention_mask) : nullptr;
+    const bool attention_defined = attention_mask &&
+        attention_mask->defined() && attention_mask->numel() > 0;
+    const bool locally_valid = input_ids && target_mask &&
+        input_ids->defined() && target_mask->defined() &&
+        input_ids->is_cuda() && target_mask->is_cuda() &&
+        input_ids->dim() == 2 && input_ids->size(1) > 1 &&
+        target_mask->sizes() == input_ids->sizes() &&
+        input_ids->scalar_type() == at::kLong &&
+        supported_mask_dtype(*target_mask) &&
+        target_mask->device() == input_ids->device() &&
+        (!attention_defined ||
+            (attention_mask->is_cuda() && attention_mask->dim() == 2 &&
+             attention_mask->sizes() == input_ids->sizes() &&
+             supported_mask_dtype(*attention_mask) &&
+             attention_mask->device() == input_ids->device()));
+
+    const std::vector<int64_t> signature{
+        locally_valid ? 1 : 0,
+        locally_valid ? input_ids->size(0) : -1,
+        locally_valid ? input_ids->size(1) : -1,
+        locally_valid ? static_cast<int64_t>(input_ids->scalar_type()) : -1,
+        locally_valid ? static_cast<int64_t>(target_mask->scalar_type()) : -1,
+        locally_valid ? static_cast<int64_t>(attention_defined) : -1,
+        locally_valid && attention_defined
+            ? static_cast<int64_t>(attention_mask->scalar_type()) : -1,
+    };
+    std::vector<int64_t> agreed_signature = signature;
+    if (window.batch_size < 0) {
+        // A PP-wide collective is safe only before the first activation enters
+        // the pipeline. Later stages intentionally run at different tick
+        // indices, so a per-tick collective would deadlock against P2P traffic.
+        auto options = at::TensorOptions().dtype(at::kLong).device(
+            at::kCUDA, ctx->cuda_device);
+        auto minimum = at::tensor(signature, options);
+        auto maximum = minimum.clone();
+        auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+        auto control_comm = ctx->pp_control_comm ?
+            ctx->pp_control_comm : ctx->pp_comm;
+        TORCH_CHECK(ncclGroupStart() == ncclSuccess,
+            "pipeline contract ncclGroupStart failed");
+        const auto min_error = ncclAllReduce(
+            minimum.data_ptr<int64_t>(), minimum.data_ptr<int64_t>(),
+            minimum.numel(), ncclInt64, ncclMin, control_comm, stream);
+        const auto max_error = ncclAllReduce(
+            maximum.data_ptr<int64_t>(), maximum.data_ptr<int64_t>(),
+            maximum.numel(), ncclInt64, ncclMax, control_comm, stream);
+        const auto end_error = ncclGroupEnd();
+        const auto error = min_error != ncclSuccess ? min_error
+            : (max_error != ncclSuccess ? max_error : end_error);
+        TORCH_CHECK(error == ncclSuccess,
+            "pipeline contract consensus failed: ", ncclGetErrorString(error));
+        const auto minimum_cpu = minimum.to(at::kCPU);
+        const auto maximum_cpu = maximum.to(at::kCPU);
+        const auto* minimum_data = minimum_cpu.data_ptr<int64_t>();
+        const auto* maximum_data = maximum_cpu.data_ptr<int64_t>();
+        for (int64_t index = 0; index < minimum.numel(); ++index) {
+            TORCH_CHECK(minimum_data[index] == maximum_data[index],
+                "pipeline tensor shape/dtype contract differs across PP stages");
+            agreed_signature[index] = minimum_data[index];
+        }
+    }
+    const auto* contract = agreed_signature.data();
+    if (window.batch_size < 0) {
+        // The first forward establishes the PP-wide tensor contract. Even when
+        // every rank reports an invalid local tensor, fail after the collective
+        // so all ranks take the same path instead of fabricating shapes.
+        TORCH_CHECK(contract[0] == 1,
+            "pipeline window forward tensors invalid on at least one PP stage");
+    } else if (contract[0] != 1) {
+        // Later local errors are converted into dummy work by the caller so
+        // activation/gradient P2P remains aligned until finish consensus.
+        return false;
+    }
+    if (window.batch_size < 0) {
+        window.batch_size = contract[1];
+        window.sequence_length = contract[2];
+        window.input_dtype = static_cast<int32_t>(contract[3]);
+        window.target_dtype = static_cast<int32_t>(contract[4]);
+        window.attention_present = static_cast<int32_t>(contract[5]);
+        window.attention_dtype = static_cast<int32_t>(contract[6]);
+    } else {
+        if (contract[1] != window.batch_size ||
+                contract[2] != window.sequence_length ||
+                contract[3] != window.input_dtype ||
+                contract[4] != window.target_dtype ||
+                contract[5] != window.attention_present ||
+                contract[6] != window.attention_dtype)
+            return false;
+    }
+    if (!std::isfinite(tick.gradient_scale) || tick.gradient_scale <= 0.0)
+        return false;
+    if (attention_defined) {
+        try {
+            validate_linear_attention_mask(ctx, *attention_mask);
+        } catch (...) {
+            return false;
+        }
+    }
+    return true;
+}
+
+static bool pipeline_window_validate_tick(
+    TrainingContext* ctx,
+    const Qwen36PipelineTickV1& tick
+) {
+    auto& window = ctx->pp_window;
+    TORCH_CHECK(window.active, "pipeline window is not active");
+    bool valid = tick.window_id == window.window_id && tick.chunk_id == 0;
+    const int32_t expected_phase = tick.forward_mb < 0 ? 2
+        : (tick.backward_mb < 0 ? 0 : 1);
+    valid = valid && tick.phase == expected_phase;
+    valid = valid && tick.forward_mb >= -1 && tick.backward_mb >= -1;
+    const int64_t warmup = std::min<int64_t>(
+        ctx->pp_world_size - ctx->pp_rank - 1,
+        window.num_microbatches);
+    const int64_t expected_forward = window.next_forward;
+    const int64_t expected_backward = window.next_backward;
+    int64_t required_forward = -1;
+    int64_t required_backward = -1;
+    if (expected_forward < warmup) {
+        required_forward = expected_forward;
+    } else if (expected_forward < window.num_microbatches) {
+        required_forward = expected_forward;
+        required_backward = expected_forward - warmup;
+    } else if (expected_backward < window.num_microbatches) {
+        required_backward = expected_backward;
+    }
+    valid = valid && tick.forward_mb == required_forward &&
+        tick.backward_mb == required_backward;
+    if (required_backward >= 0) {
+        auto it = window.slots.find(required_backward);
+        const bool same_tick_forward = required_forward == required_backward &&
+            required_forward >= 0;
+        if (!same_tick_forward) {
+            valid = valid && it != window.slots.end() &&
+                it->second.forward_done && !it->second.backward_done;
+        }
+    }
+    return valid;
+}
+
+static at::Tensor pipeline_window_forward(
+    TrainingContext* ctx,
+    const Qwen36PipelineTickV1& tick,
+    int64_t hidden_size,
+    const at::Tensor& received_stage_input = at::Tensor()
+) {
+    auto& input_ids = *reinterpret_cast<at::Tensor*>(tick.input_ids);
+    auto& target_mask = *reinterpret_cast<at::Tensor*>(tick.target_mask);
+    auto* attention_mask = tick.attention_mask
+        ? reinterpret_cast<at::Tensor*>(tick.attention_mask) : nullptr;
+    if (attention_mask && attention_mask->defined() &&
+            attention_mask->numel() > 0) {
+        ctx->attention_mask = *attention_mask;
+        elide_trivial_attention_mask(ctx);
+    } else if (ctx->pad_token_id >= 0) {
+        ctx->attention_mask = derive_attention_mask(ctx, input_ids);
+        elide_trivial_attention_mask(ctx);
+    } else {
+        ctx->attention_mask = at::Tensor();
+        ctx->attention_lengths = at::Tensor();
+    }
+    const double supervised_tokens = target_mask.narrow(
+        1, 1, target_mask.size(1) - 1).to(at::kFloat).sum().item<double>();
+    const double token_weight = tick.gradient_scale * supervised_tokens;
+    TORCH_CHECK(std::isfinite(token_weight) && token_weight >= 0.0,
+        "pipeline window token weight must be finite and non-negative");
+    TORCH_CHECK(ctx->pp_window.slots.find(tick.forward_mb) ==
+            ctx->pp_window.slots.end(),
+        "pipeline forward microbatch slot is already live");
+
+    PipelineWindowSlot slot;
+    slot.input_ids = input_ids;
+    slot.target_mask = target_mask;
+    if (ctx->attention_mask.defined()) slot.attention_mask = ctx->attention_mask;
+    slot.gradient_scale = tick.gradient_scale;
+    slot.token_weight = token_weight;
+    slot.row_token_weights = target_mask.narrow(
+        1, 1, target_mask.size(1) - 1).to(at::kFloat).sum(1)
+        .mul(tick.gradient_scale).contiguous();
+    // Each in-flight microbatch gets an independent activation graph. Fixed
+    // LoRA parameters and their singleton views are immutable until finish,
+    // so keep the projection metadata cache for the whole pipeline window.
+    ctx->lora_cache_valid = false;
+    if (ctx->pp_window.dynamic_lora) {
+        ctx->lora_batch_valid = true;
+        prepare_lora_batch(ctx, 0, -1,
+            &ctx->pp_window.selected_adapter_indices);
+    } else if (!ctx->lora_batch_valid) {
+        prepare_fixed_lora_batch(ctx);
+    }
+    at::AutoGradMode grad_enable(true);
+    if (ctx->is_first_pipeline_stage) {
+        TORCH_CHECK(!received_stage_input.defined(),
+            "first pipeline stage must not receive an activation tensor");
+        slot.stage_input = vocabulary_embedding(ctx, input_ids);
+    } else {
+        TORCH_CHECK(received_stage_input.defined() &&
+                received_stage_input.is_cuda() &&
+                received_stage_input.scalar_type() == ctx->compute_type &&
+                received_stage_input.dim() == 3 &&
+                received_stage_input.size(0) == input_ids.size(0) &&
+                received_stage_input.size(1) == input_ids.size(1) &&
+                received_stage_input.size(2) == hidden_size,
+            "pipeline stage received an activation outside the fixed contract");
+        slot.stage_input = received_stage_input;
+        slot.stage_input.set_requires_grad(true);
+    }
+    slot.stage_output = forward_stage_layers(ctx, slot.stage_input);
+    slot.forward_done = true;
+    auto inserted = ctx->pp_window.slots.emplace(tick.forward_mb, std::move(slot));
+    TORCH_CHECK(inserted.second, "pipeline forward slot insertion failed");
+    auto& live_slot = inserted.first->second;
+    if (ctx->pp_window.dynamic_lora)
+        ctx->pp_window.adapter_token_counts.add_(live_slot.row_token_weights);
+    if (!ctx->pp_window.normalization_mask.defined())
+        ctx->pp_window.normalization_mask = target_mask;
+    ctx->pp_window.next_forward++;
+    return live_slot.stage_output;
+}
+
+static at::Tensor pipeline_window_backward(
+    TrainingContext* ctx,
+    int64_t microbatch_id,
+    const at::Tensor& received_output_grad = at::Tensor()
+) {
+    auto& window = ctx->pp_window;
+    auto it = window.slots.find(microbatch_id);
+    TORCH_CHECK(it != window.slots.end() && it->second.forward_done &&
+            !it->second.backward_done,
+        "pipeline backward slot is not ready");
+    auto& slot = it->second;
+    if (slot.attention_mask.defined()) {
+        validate_linear_attention_mask(ctx, slot.attention_mask);
+        ctx->attention_mask = slot.attention_mask;
+        elide_trivial_attention_mask(ctx);
+    } else {
+        ctx->attention_mask = at::Tensor();
+        ctx->attention_lengths = at::Tensor();
+    }
+    double local_loss = 0.0;
+    double local_loss_numerator = 0.0;
+    double local_loss_denominator = 0.0;
+    if (ctx->is_last_pipeline_stage) {
+        TORCH_CHECK(!received_output_grad.defined(),
+            "last pipeline stage must not receive an output gradient");
+        auto loss = pipeline_compute_loss(
+            ctx, slot.stage_output, slot.input_ids, slot.target_mask);
+        local_loss = loss.value;
+        local_loss_numerator = loss.numerator * slot.gradient_scale;
+        local_loss_denominator = loss.denominator * slot.gradient_scale;
+        if (window.dynamic_lora) {
+            TORCH_CHECK(loss.sample_loss_numerators.defined() &&
+                    loss.sample_token_counts.defined() &&
+                    loss.sample_loss_numerators.numel() ==
+                        window.loss_numerators.numel() &&
+                    loss.sample_token_counts.numel() ==
+                        window.loss_token_counts.numel(),
+                "dynamic pipeline loss statistics do not match selected adapters");
+            window.loss_numerators.add_(
+                loss.sample_loss_numerators * slot.gradient_scale);
+            window.loss_token_counts.add_(
+                loss.sample_token_counts * slot.gradient_scale);
+        }
+        auto hidden_grad = window.dynamic_lora
+            ? loss.hidden_grad * slot.row_token_weights.reshape({-1, 1, 1})
+            : loss.hidden_grad * slot.token_weight;
+        slot.stage_output.backward(hidden_grad);
+    } else {
+        TORCH_CHECK(received_output_grad.defined() &&
+                received_output_grad.sizes() == slot.stage_output.sizes() &&
+                received_output_grad.scalar_type() == slot.stage_output.scalar_type(),
+            "pipeline output gradient violates the fixed activation contract");
+        slot.stage_output.backward(received_output_grad);
+    }
+    at::Tensor input_grad;
+    if (!ctx->is_first_pipeline_stage) {
+        input_grad = slot.stage_input.grad();
+        TORCH_CHECK(input_grad.defined(),
+            "pipeline window stage did not produce input gradient");
+        input_grad = input_grad.contiguous();
+    }
+    harvest_gradient_accumulators(ctx);
+    if (window.dynamic_lora) {
+        window.total_loss += local_loss_numerator;
+        window.total_loss_weight += local_loss_denominator;
+    } else {
+        window.total_loss += local_loss;
+    }
+    if (!window.dynamic_lora) ctx->accumulated_token_weight += slot.token_weight;
+    slot.backward_done = true;
+    window.next_backward++;
+    window.slots.erase(it);
+    ctx->lora_cache_valid = false;
+    return input_grad;
+}
+
+static void pipeline_window_reset(TrainingContext* ctx) {
+    if (!ctx) return;
+    if (ctx->pp_window.padding_mode_overridden)
+        ctx->pad_heterogeneous_lora_batch =
+            ctx->pp_window.previous_pad_heterogeneous_lora_batch;
+    ctx->pp_window = PipelineWindowState{};
+}
+
+extern "C" __attribute__((visibility("default"))) int32_t qwen36_pipeline_begin_v1(
+    void* ctx_ptr,
+    const Qwen36PipelineWindowV1* window_spec
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "pipeline begin requires a context");
+        TORCH_CHECK(ctx->pp_world_size >= 2 &&
+                (ctx->pp_control_comm || ctx->pp_comm),
+            "pipeline window requires a PP communicator with PP_SIZE >= 2");
+        const bool valid_abi = window_spec &&
+            window_spec->struct_size >= sizeof(Qwen36PipelineWindowV1) &&
+            window_spec->version == 1;
+        const bool dynamic_lora = valid_abi &&
+            window_spec->flags == kPipelineWindowFlagDynamicLora;
+        const bool supported_flags = valid_abi &&
+            (window_spec->flags == 0 || dynamic_lora);
+        const bool local_valid = valid_abi && supported_flags &&
+            !ctx->pp_window.active && ctx->cp_world_size == 1 &&
+            ctx->router_aux_loss_coef == 0.0 &&
+            window_spec->window_id >= 0 &&
+            window_spec->num_microbatches > 0 &&
+            window_spec->num_microbatches <= 4096 &&
+            window_spec->schedule == 0 && window_spec->num_chunks == 1 &&
+            (dynamic_lora ? !ctx->adapters.empty() : ctx->adapters.empty()) &&
+            !ctx->has_mtp && !ctx->use_checkpoint &&
+            !ctx->accumulation_active &&
+            ctx->accumulated_token_weight == 0.0 &&
+            gdn_state_checkpoint_environment_valid();
+        const bool globally_valid = pipeline_control_all_succeeded(
+            ctx, local_valid);
+        TORCH_CHECK(local_valid && globally_valid,
+            "pipeline begin preflight failed on one or more PP stages");
+        validate_pipeline_window_collective_registry(
+            ctx, window_spec->window_id, window_spec->num_microbatches,
+            window_spec->schedule, window_spec->num_chunks, window_spec->flags);
+        ctx->pp_window.active = true;
+        ctx->pp_window.window_id = window_spec->window_id;
+        ctx->pp_window.num_microbatches = window_spec->num_microbatches;
+        ctx->pp_window.schedule = window_spec->schedule;
+        ctx->pp_window.num_chunks = window_spec->num_chunks;
+        ctx->pp_window.flags = window_spec->flags;
+        ctx->pp_window.dynamic_lora = dynamic_lora;
+        if (dynamic_lora) {
+            // PP keeps one projection cache alive across the whole 1F1B
+            // window.  Always use the padded representation here so a
+            // selected request may mix ranks and target-module holes without
+            // discovering a stack/shape mismatch after P2P has started.
+            ctx->pp_window.previous_pad_heterogeneous_lora_batch =
+                ctx->pad_heterogeneous_lora_batch;
+            ctx->pp_window.padding_mode_overridden = true;
+            ctx->pad_heterogeneous_lora_batch = true;
+        }
+        if (dynamic_lora) {
+            ctx->pp_window.selected_adapter_indices.clear();
+            for (size_t index = 0; index < ctx->adapters.size(); ++index)
+                ctx->pp_window.selected_adapter_indices.push_back(index);
+            ctx->pp_window.adapter_token_counts = at::zeros(
+                {static_cast<int64_t>(ctx->adapters.size())},
+                at::TensorOptions().dtype(at::kFloat).device(
+                    at::kCUDA, ctx->cuda_device));
+            ctx->pp_window.loss_numerators = at::zeros_like(
+                ctx->pp_window.adapter_token_counts);
+            ctx->pp_window.loss_token_counts = at::zeros_like(
+                ctx->pp_window.adapter_token_counts);
+        }
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[pipeline_window] begin FAILED: %s\n", e.what());
+        return -1;
+    } catch (...) {
+        fprintf(stderr, "[pipeline_window] begin FAILED: unknown error\n");
+        return -1;
+    }
+}
+
+// Selected dynamic tenants use the same v1 tick/finish ABI. The selection is
+// validated before opening the window, then stored as stable registry indices;
+// unselected adapters remain untouched by the dynamic finalizer.
+extern "C" __attribute__((visibility("default"))) int32_t
+qwen36_pipeline_begin_dynamic_selected_v1(
+    void* ctx_ptr,
+    const Qwen36PipelineWindowV1* window_spec,
+    const int64_t* adapter_ids,
+    int32_t adapter_count
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        const bool local_valid = ctx && adapter_ids && adapter_count > 0 &&
+            adapter_count <= static_cast<int32_t>(ctx->adapters.size()) &&
+            window_spec && window_spec->flags == kPipelineWindowFlagDynamicLora;
+        std::vector<size_t> selected;
+        std::vector<int64_t> ids;
+        if (local_valid) {
+            ids.assign(adapter_ids, adapter_ids + adapter_count);
+            for (const auto id : ids) {
+                auto it = std::find_if(ctx->adapters.begin(), ctx->adapters.end(),
+                    [id](const auto& adapter) { return adapter.id == id; });
+                if (id <= 0 || it == ctx->adapters.end()) {
+                    selected.clear();
+                    break;
+                }
+                selected.push_back(static_cast<size_t>(
+                    std::distance(ctx->adapters.begin(), it)));
+            }
+            if (selected.size() != ids.size()) selected.clear();
+            if (!selected.empty()) {
+                std::set<size_t> unique(selected.begin(), selected.end());
+                if (unique.size() != selected.size()) selected.clear();
+            }
+        }
+        const bool globally_valid = adapter_collective_all_succeeded(
+            ctx, local_valid && !selected.empty());
+        TORCH_CHECK(local_valid && !selected.empty() && globally_valid,
+            "selected dynamic pipeline adapter request is invalid on one or more ranks");
+        validate_pipeline_selected_adapter_request(
+            ctx, adapter_ids, adapter_count);
+        const int32_t status = qwen36_pipeline_begin_v1(ctx_ptr, window_spec);
+        if (status != 0) return status;
+        ctx->pp_window.selected_adapter_indices = std::move(selected);
+        ctx->pp_window.adapter_token_counts = at::zeros(
+            {adapter_count}, at::TensorOptions().dtype(at::kFloat).device(
+                at::kCUDA, ctx->cuda_device));
+        ctx->pp_window.loss_numerators = at::zeros_like(
+            ctx->pp_window.adapter_token_counts);
+        ctx->pp_window.loss_token_counts = at::zeros_like(
+            ctx->pp_window.adapter_token_counts);
+        return 0;
+    } catch (const std::exception& error) {
+        fprintf(stderr, "[pipeline_window] selected begin FAILED: %s\n",
+            error.what());
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        if (ctx && ctx->pp_window.active) pipeline_window_reset(ctx);
+        return -1;
+    } catch (...) {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        if (ctx && ctx->pp_window.active) pipeline_window_reset(ctx);
+        return -1;
+    }
+}
+
+extern "C" __attribute__((visibility("default"))) int32_t qwen36_pipeline_tick_v1(
+    void* ctx_ptr,
+    const Qwen36PipelineTickV1* tick,
+    Qwen36PipelineResultV1* result
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx && tick,
+            "pipeline tick requires a context and tick specification");
+        TORCH_CHECK(tick->struct_size >= sizeof(Qwen36PipelineTickV1) &&
+                tick->version == 1,
+            "unsupported pipeline tick ABI");
+        pipeline_window_validate_result(result);
+        auto& window = ctx->pp_window;
+        const bool tick_valid = pipeline_window_validate_tick(ctx, *tick);
+        if (!tick_valid) window.local_error = true;
+        const int64_t warmup = std::min<int64_t>(
+            ctx->pp_world_size - ctx->pp_rank - 1,
+            window.num_microbatches);
+        const int64_t expected_forward = window.next_forward;
+        const int64_t expected_backward = window.next_backward;
+        int64_t required_forward = -1;
+        int64_t required_backward = -1;
+        if (expected_forward < warmup) {
+            required_forward = expected_forward;
+        } else if (expected_forward < window.num_microbatches) {
+            required_forward = expected_forward;
+            required_backward = expected_forward - warmup;
+        } else if (expected_backward < window.num_microbatches) {
+            required_backward = expected_backward;
+        }
+        const int64_t hidden_size = ctx->weight_ptrs.empty()
+            ? 0 : ctx->weight_ptrs[0]->numel();
+        TORCH_CHECK(hidden_size > 0,
+            "pipeline stage has no hidden-size metadata");
+        const bool is_first_stage = ctx->is_first_pipeline_stage;
+        const bool is_last_stage = ctx->is_last_pipeline_stage;
+        const bool has_forward = required_forward >= 0;
+        const bool has_backward = required_backward >= 0;
+        if (has_forward) {
+            Qwen36PipelineTickV1 effective_tick = *tick;
+            effective_tick.window_id = window.window_id;
+            effective_tick.forward_mb = required_forward;
+            effective_tick.backward_mb = required_backward;
+            effective_tick.chunk_id = 0;
+            effective_tick.phase = has_backward ? 1 : 0;
+            // Always enter the first contract handshake on canonical forwards.
+            // Short-circuiting this call on a local metadata error would let
+            // one PP rank skip the control collective while others wait.
+            bool forward_valid = tick_valid;
+            at::Tensor fallback_input_ids;
+            at::Tensor fallback_target_mask;
+            at::Tensor fallback_attention_mask;
+            at::Tensor dynamic_input_ids;
+            at::Tensor dynamic_target_mask;
+            at::Tensor dynamic_attention_mask;
+            if (window.dynamic_lora && tick->input_ids && tick->target_mask) {
+                auto& raw_ids = *reinterpret_cast<at::Tensor*>(tick->input_ids);
+                auto& raw_targets = *reinterpret_cast<at::Tensor*>(tick->target_mask);
+                const int64_t tenants = static_cast<int64_t>(
+                    window.selected_adapter_indices.size());
+                if (raw_ids.defined() && raw_targets.defined() &&
+                        raw_ids.dim() == 2 && raw_targets.dim() == 2 &&
+                        raw_ids.size(0) != 1 && raw_ids.size(0) != tenants)
+                    forward_valid = false;
+                if (raw_ids.defined() && raw_targets.defined() &&
+                        raw_ids.dim() == 2 && raw_targets.dim() == 2 &&
+                        raw_ids.size(0) == 1 && raw_targets.size(0) == 1) {
+                    dynamic_input_ids = raw_ids.repeat({tenants, 1});
+                    dynamic_target_mask = raw_targets.repeat({tenants, 1});
+                    if (tick->attention_mask) {
+                        auto& raw_attention = *reinterpret_cast<at::Tensor*>(
+                            tick->attention_mask);
+                        if (raw_attention.defined() && raw_attention.dim() == 2 &&
+                                raw_attention.size(0) == 1)
+                            dynamic_attention_mask = raw_attention.repeat({tenants, 1});
+                    }
+                } else {
+                    dynamic_input_ids = raw_ids;
+                    dynamic_target_mask = raw_targets;
+                    if (tick->attention_mask)
+                        dynamic_attention_mask = *reinterpret_cast<at::Tensor*>(
+                            tick->attention_mask);
+                }
+                if (dynamic_input_ids.defined() && dynamic_target_mask.defined()) {
+                    effective_tick.input_ids = &dynamic_input_ids;
+                    effective_tick.target_mask = &dynamic_target_mask;
+                    effective_tick.attention_mask = dynamic_attention_mask.defined()
+                        ? &dynamic_attention_mask : nullptr;
+                }
+            }
+            const bool forward_contract_valid =
+                pipeline_window_validate_forward_contract(ctx, effective_tick);
+            forward_valid = forward_valid && forward_contract_valid;
+            if (!forward_valid) {
+                window.local_error = true;
+                int64_t fallback_batch = window.batch_size;
+                int64_t fallback_sequence = window.sequence_length;
+                if (tick->input_ids) {
+                    const auto& candidate = *reinterpret_cast<at::Tensor*>(
+                        tick->input_ids);
+                    if (candidate.defined() && candidate.dim() == 2) {
+                        if (fallback_batch <= 0) fallback_batch = candidate.size(0);
+                        if (fallback_sequence <= 1)
+                            fallback_sequence = candidate.size(1);
+                    }
+                }
+                if (window.dynamic_lora && fallback_batch > 0)
+                    fallback_batch = static_cast<int64_t>(
+                        window.selected_adapter_indices.size());
+                if ((fallback_batch <= 0 || fallback_sequence <= 1) &&
+                        tick->target_mask) {
+                    const auto& candidate = *reinterpret_cast<at::Tensor*>(
+                        tick->target_mask);
+                    if (candidate.defined() && candidate.dim() == 2) {
+                        if (fallback_batch <= 0) fallback_batch = candidate.size(0);
+                        if (fallback_sequence <= 1)
+                            fallback_sequence = candidate.size(1);
+                    }
+                }
+                TORCH_CHECK(fallback_batch > 0 && fallback_sequence > 1,
+                    "pipeline window cannot recover a malformed first input");
+                std::vector<int64_t> shape_values{
+                    fallback_batch, fallback_sequence};
+                auto shape = at::IntArrayRef(shape_values);
+                auto device = at::Device(at::kCUDA, ctx->cuda_device);
+                fallback_input_ids = at::zeros(
+                    shape, at::TensorOptions().device(device).dtype(at::kLong));
+                fallback_target_mask = at::zeros(
+                    shape, at::TensorOptions().device(device).dtype(at::kLong));
+                if (window.attention_present > 0)
+                    fallback_attention_mask = at::ones(
+                        shape, at::TensorOptions().device(device).dtype(at::kFloat));
+                effective_tick.input_ids = &fallback_input_ids;
+                effective_tick.target_mask = &fallback_target_mask;
+                effective_tick.attention_mask =
+                    fallback_attention_mask.defined() ? &fallback_attention_mask : nullptr;
+                effective_tick.gradient_scale = 1.0;
+            }
+            auto& input_ids = *reinterpret_cast<at::Tensor*>(
+                effective_tick.input_ids);
+            at::Tensor stage_input;
+            if (!is_first_stage) {
+                stage_input = pipeline_recv_tensor(
+                    ctx, input_ids, ctx->pp_rank - 1, hidden_size);
+            }
+            auto stage_output = pipeline_window_forward(
+                ctx, effective_tick, hidden_size, stage_input).contiguous();
+
+            if (is_last_stage) {
+                auto input_grad = pipeline_window_backward(
+                    ctx, required_backward).contiguous();
+                if (!is_first_stage)
+                    pipeline_send_tensor(ctx, input_grad, ctx->pp_rank - 1);
+            } else if (has_backward) {
+                auto& backward_slot = window.slots.at(required_backward);
+                auto output_grad = pipeline_exchange_tensor(
+                    ctx, stage_output, backward_slot.input_ids,
+                    ctx->pp_rank + 1, hidden_size);
+                auto input_grad = pipeline_window_backward(
+                    ctx, required_backward, output_grad);
+                if (!is_first_stage)
+                    pipeline_send_tensor(ctx, input_grad, ctx->pp_rank - 1);
+            } else {
+                pipeline_send_tensor(ctx, stage_output, ctx->pp_rank + 1);
+            }
+        } else {
+            auto& backward_slot = window.slots.at(required_backward);
+            at::Tensor output_grad;
+            if (!is_last_stage) {
+                output_grad = pipeline_recv_tensor(
+                    ctx, backward_slot.input_ids,
+                    ctx->pp_rank + 1, hidden_size);
+            }
+            auto input_grad = pipeline_window_backward(
+                ctx, required_backward, output_grad);
+            if (!is_first_stage)
+                pipeline_send_tensor(ctx, input_grad, ctx->pp_rank - 1);
+        }
+        pipeline_window_write_result(ctx, result, 0, window.total_loss);
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[pipeline_window] tick FAILED: %s\n", e.what());
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        clear_gradient_accumulators(ctx);
+        pipeline_window_reset(ctx);
+        return -1;
+    } catch (...) {
+        fprintf(stderr, "[pipeline_window] tick FAILED: unknown error\n");
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        clear_gradient_accumulators(ctx);
+        pipeline_window_reset(ctx);
+        return -1;
+    }
+}
+
+extern "C" __attribute__((visibility("default"))) int32_t qwen36_pipeline_finish_v1(
+    void* ctx_ptr,
+    int32_t apply_optimizer,
+    Qwen36PipelineResultV1* result
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx && ctx->pp_window.active,
+            "pipeline finish requires an active window");
+        pipeline_window_validate_result(result);
+        TORCH_CHECK(apply_optimizer == 1,
+            "pipeline window finish must apply exactly one optimizer update");
+        auto& window = ctx->pp_window;
+        TORCH_CHECK(window.next_forward == window.num_microbatches &&
+                window.next_backward == window.num_microbatches &&
+                window.slots.empty() &&
+                !window.pending_backward_send.defined() &&
+                window.pending_backward_mb == -1,
+            "pipeline window cannot finish before every forward/backward completes");
+        TORCH_CHECK(window.normalization_mask.defined(),
+            "pipeline window has no normalization mask");
+        auto status = at::full(
+            {1}, window.local_error ? 0 : 1,
+            at::TensorOptions().device(at::kCUDA, ctx->cuda_device).dtype(at::kInt));
+        auto control_comm = ctx->pp_control_comm ?
+            ctx->pp_control_comm : ctx->pp_comm;
+        auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+        const auto status_error = ncclAllReduce(
+            status.data_ptr<int32_t>(), status.data_ptr<int32_t>(), 1,
+            ncclInt32, ncclMin, control_comm, stream);
+        TORCH_CHECK(status_error == ncclSuccess,
+            "pipeline error consensus failed: ",
+            ncclGetErrorString(status_error));
+        TORCH_CHECK(status.to(at::kCPU).item<int32_t>() == 1,
+            "pipeline window aborted because at least one PP rank rejected a tick");
+        if (window.dynamic_lora) {
+            TORCH_CHECK(window.adapter_token_counts.defined(),
+                "dynamic pipeline window has no adapter token counts");
+            // Every PP stage receives the same tenant rows. Compare the
+            // accumulated token vector before the optimizer boundary so a
+            // rank-local mask/value divergence cannot silently skew only one
+            // stage's numerator normalization.
+            auto token_min = window.adapter_token_counts.clone();
+            auto token_max = window.adapter_token_counts.clone();
+            auto token_stream = c10::cuda::getCurrentCUDAStream(
+                ctx->cuda_device).stream();
+            auto token_comm = ctx->pp_control_comm ?
+                ctx->pp_control_comm : ctx->pp_comm;
+            auto token_min_error = ncclAllReduce(
+                window.adapter_token_counts.data_ptr<float>(),
+                token_min.data_ptr<float>(), token_min.numel(), ncclFloat,
+                ncclMin, token_comm, token_stream);
+            auto token_max_error = ncclAllReduce(
+                window.adapter_token_counts.data_ptr<float>(),
+                token_max.data_ptr<float>(), token_max.numel(), ncclFloat,
+                ncclMax, token_comm, token_stream);
+            TORCH_CHECK(token_min_error == ncclSuccess &&
+                    token_max_error == ncclSuccess,
+                "dynamic pipeline token-count consensus failed");
+            TORCH_CHECK(at::equal(token_min, token_max),
+                "dynamic pipeline token counts differ across PP stages");
+            auto full_adapter_token_counts = at::zeros(
+                {static_cast<int64_t>(ctx->adapters.size())},
+                window.adapter_token_counts.options());
+            for (size_t ordinal = 0;
+                 ordinal < window.selected_adapter_indices.size(); ++ordinal) {
+                full_adapter_token_counts.select(
+                    0, static_cast<int64_t>(window.selected_adapter_indices[ordinal]))
+                    .copy_(window.adapter_token_counts.select(
+                        0, static_cast<int64_t>(ordinal)));
+            }
+            int32_t finalizer_phase = 0;
+            int64_t max_rank = 1;
+            for (const auto& adapter : ctx->adapters)
+                max_rank = std::max(max_rank, adapter.rank);
+            auto dummy_ids = at::zeros(
+                {static_cast<int64_t>(ctx->adapters.size()), 2},
+                at::TensorOptions().dtype(at::kLong).device(
+                    at::kCUDA, ctx->cuda_device));
+            auto dummy_mask = at::zeros_like(dummy_ids);
+            const double finalizer_loss = qwen36_train_multi_lora_impl(
+                ctx, &dummy_ids, &dummy_mask, nullptr,
+                static_cast<int32_t>(ctx->adapters.size()),
+                static_cast<int32_t>(max_rank),
+                static_cast<DynamicMultiLoraMode>(2), &finalizer_phase,
+                nullptr, nullptr, true, &full_adapter_token_counts, true,
+                false);
+            TORCH_CHECK(finalizer_loss >= 0.0 && finalizer_phase >= 1,
+                "dynamic pipeline finalizer failed");
+            const int64_t completed_microbatches = window.num_microbatches;
+            const double aggregate_loss = window.total_loss /
+                std::max(1.0, window.total_loss_weight);
+            const double loss = pipeline_broadcast_loss(
+                ctx, aggregate_loss, window.normalization_mask.options());
+            TORCH_CHECK(std::isfinite(loss) && loss >= 0.0,
+                "pipeline dynamic-LoRA loss became non-finite");
+            pipeline_window_write_result(ctx, result, 0, loss);
+            pipeline_window_reset(ctx);
+            if (result) {
+                result->completed_fwd = completed_microbatches;
+                result->completed_bwd = completed_microbatches;
+                result->in_flight = 0;
+            }
+            return 0;
+        }
+        validate_pipeline_collective_registry(
+            ctx, 1, 0.0, ctx->accumulated_token_weight);
+        const int64_t completed_microbatches = window.num_microbatches;
+        const double aggregate_loss = window.total_loss /
+            static_cast<double>(completed_microbatches);
+        const double loss = pipeline_broadcast_loss(
+            ctx, aggregate_loss, window.normalization_mask.options());
+        TORCH_CHECK(std::isfinite(loss) && loss >= 0.0,
+            "pipeline fixed-LoRA loss became non-finite");
+        apply_fixed_lora_optimizer(
+            ctx, window.normalization_mask, ctx->accumulated_token_weight);
+        pipeline_window_write_result(ctx, result, 0, loss);
+        pipeline_window_reset(ctx);
+        if (result) {
+            result->completed_fwd = completed_microbatches;
+            result->completed_bwd = completed_microbatches;
+            result->in_flight = 0;
+        }
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[pipeline_window] finish FAILED: %s\n", e.what());
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        clear_gradient_accumulators(ctx);
+        pipeline_window_reset(ctx);
+        return -1;
+    } catch (...) {
+        fprintf(stderr, "[pipeline_window] finish FAILED: unknown error\n");
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        clear_gradient_accumulators(ctx);
+        pipeline_window_reset(ctx);
+        return -1;
+    }
+}
+
+extern "C" __attribute__((visibility("default"))) int32_t qwen36_pipeline_abort_v1(
+    void* ctx_ptr
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "pipeline abort requires a context");
+        clear_gradient_accumulators(ctx);
+        pipeline_window_reset(ctx);
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[pipeline_window] abort FAILED: %s\n", e.what());
+        return -1;
+    } catch (...) {
+        fprintf(stderr, "[pipeline_window] abort FAILED: unknown error\n");
+        return -1;
+    }
+}
+
+// One synchronous pipeline microbatch. This is intentionally a correctness
+// baseline: each rank owns a contiguous stage and exchanges one activation and
+// one gradient per step. It can later be replaced by a 1F1B scheduler without
+// changing the stage-local layer or optimizer boundaries.
+static double qwen36_pipeline_train_micro_step(
+    TrainingContext* ctx,
+    at::Tensor& input_ids,
+    at::Tensor& target_mask,
+    at::Tensor* attention_mask,
+    double gradient_scale,
+    int32_t apply_optimizer
+) {
+    TORCH_CHECK(ctx && ctx->pp_world_size > 1 && ctx->cp_world_size == 1,
+        "pipeline execution requires PP_SIZE>1 and CP_SIZE=1");
+    TORCH_CHECK(ctx->adapters.empty(),
+        "pipeline fixed-LoRA path does not yet support dynamic adapters");
+    TORCH_CHECK(ctx->router_aux_loss_coef == 0.0,
+        "pipeline router auxiliary loss is not yet supported");
+    TORCH_CHECK(!ctx->has_mtp && !ctx->use_checkpoint,
+        "pipeline baseline does not support MTP or manual checkpointing");
+    TORCH_CHECK(apply_optimizer == 0 || apply_optimizer == 1,
+        "pipeline apply_optimizer must be 0 or 1");
+    TORCH_CHECK(std::isfinite(gradient_scale) && gradient_scale > 0.0,
+        "pipeline gradient scale must be finite and positive");
+    TORCH_CHECK(input_ids.is_cuda() && target_mask.is_cuda() &&
+            input_ids.dim() == 2 && target_mask.sizes() == input_ids.sizes() &&
+            input_ids.scalar_type() == at::kLong,
+        "invalid pipeline input tensors");
+    if (attention_mask) {
+        validate_linear_attention_mask(ctx, *attention_mask);
+        ctx->attention_mask = *attention_mask;
+        elide_trivial_attention_mask(ctx);
+    } else if (ctx->pad_token_id >= 0) {
+        ctx->attention_mask = derive_attention_mask(ctx, input_ids);
+        elide_trivial_attention_mask(ctx);
+    } else {
+        ctx->attention_mask = at::Tensor();
+        ctx->attention_lengths = at::Tensor();
+    }
+    const double supervised_tokens = target_mask.narrow(
+        1, 1, target_mask.size(1) - 1).to(at::kFloat).sum().item<double>();
+    const double micro_token_weight = gradient_scale * supervised_tokens;
+    const double next_accumulated_token_weight =
+        ctx->accumulated_token_weight + micro_token_weight;
+    TORCH_CHECK(std::isfinite(micro_token_weight) && micro_token_weight >= 0.0,
+        "pipeline token weight must be finite and non-negative");
+    TORCH_CHECK(std::isfinite(next_accumulated_token_weight) &&
+            next_accumulated_token_weight >= 0.0,
+        "pipeline accumulated token weight must be finite and non-negative");
+    validate_pipeline_collective_registry(
+        ctx, apply_optimizer, micro_token_weight, next_accumulated_token_weight);
+    const int64_t hidden_size = ctx->weight_ptrs.empty()
+        ? 0 : ctx->weight_ptrs[0]->numel();
+    TORCH_CHECK(hidden_size > 0, "pipeline stage has no hidden-size metadata");
+
+    // Pipeline fixed-LoRA execution must retain the A/B graph for backward
+    // even when TP=1; the legacy materialized delta cache is inference-only.
+    prepare_fixed_lora_batch(ctx);
+    at::AutoGradMode grad_enable(true);
+    at::Tensor stage_input;
+    if (ctx->is_first_pipeline_stage) {
+        stage_input = vocabulary_embedding(ctx, input_ids);
+    } else {
+        stage_input = pipeline_recv_tensor(
+            ctx, input_ids, ctx->pp_rank - 1, hidden_size);
+        stage_input.set_requires_grad(true);
+    }
+    auto stage_output = forward_stage_layers(ctx, stage_input);
+
+    double local_loss = 0.0;
+    if (ctx->is_last_pipeline_stage) {
+        auto loss = pipeline_compute_loss(
+            ctx, stage_output, input_ids, target_mask);
+        local_loss = loss.value;
+        auto hidden_grad = loss.hidden_grad * micro_token_weight;
+        stage_output.backward(hidden_grad);
+        if (!ctx->is_first_pipeline_stage) {
+            auto input_grad = stage_input.grad();
+            TORCH_CHECK(input_grad.defined(),
+                "pipeline last stage did not produce input gradient");
+            pipeline_send_tensor(ctx, input_grad, ctx->pp_rank - 1);
+        }
+    } else {
+        pipeline_send_tensor(ctx, stage_output, ctx->pp_rank + 1);
+        auto output_grad = pipeline_recv_tensor(
+            ctx, input_ids, ctx->pp_rank + 1, hidden_size);
+        stage_output.backward(output_grad);
+        if (!ctx->is_first_pipeline_stage) {
+            auto input_grad = stage_input.grad();
+            TORCH_CHECK(input_grad.defined(),
+                "pipeline stage did not produce input gradient");
+            pipeline_send_tensor(ctx, input_grad, ctx->pp_rank - 1);
+        }
+    }
+
+    harvest_gradient_accumulators(ctx);
+    ctx->accumulated_token_weight = next_accumulated_token_weight;
+    const double loss = pipeline_broadcast_loss(
+        ctx, local_loss, input_ids.options());
+    TORCH_CHECK(std::isfinite(loss) && loss >= 0.0,
+        "pipeline fixed-LoRA loss became non-finite");
+    if (apply_optimizer) {
+        apply_fixed_lora_optimizer(
+            ctx, target_mask, ctx->accumulated_token_weight);
+    } else {
+        const bool local_gradients_finite =
+            fixed_lora_accumulators_are_finite(ctx);
+        TORCH_CHECK(fixed_optimizer_collective_all_succeeded(
+                ctx, local_gradients_finite) && local_gradients_finite,
+            "pipeline fixed-LoRA accumulated gradient became non-finite");
+    }
+    ctx->lora_cache_valid = false;
+    ctx->lora_batch_valid = false;
+    return loss;
+}
+
+// Manual sequential backward — recompute each group with grad, backprop, free.
+// Only 1 group's intermediate tensors exist at any time.
+static void manual_group_backward(
+    TrainingContext* ctx,
+    const at::Tensor& hidden_grad
+) {
+    auto& groups = ctx->group_ranges;
+    int64_t num_groups = (int64_t)groups.size();
+    at::Tensor grad = hidden_grad;
+    at::AutoGradMode grad_mode(true);
+
+    for (int64_t g = num_groups - 1; g >= 0; g--) {
+        int64_t start = groups[g].first;
+        int64_t end = groups[g].second;
+
+        // Restore input from saved (CPU if offloaded)
+        auto input = ctx->group_inputs[g].to(hidden_grad.device()).detach().set_requires_grad(true);
+
+        // Each recomputed group builds a fresh autograd graph. Reusing a
+        // differentiable LoRA cache after grad(..., retain_graph=false) would
+        // point the next group at freed graph nodes.
+        ctx->lora_batch_valid = false;
+        ctx->lora_cache_valid = false;
+        if (!ctx->adapters.empty()) prepare_lora_batch(ctx, start, end);
+        else if (ctx->tp_world_size > 1 || ctx->cp_world_size > 1)
+            prepare_fixed_lora_batch(ctx);
+        else precompute_lora_cache(ctx);
+
+        // Recompute forward with grad for this group only
+        auto output = forward_layer_group(ctx, input, start, end);
+
+        // Backprop through this group using grad() instead of backward().
+        // grad() only computes gradients for specified inputs — faster than
+        // backward() which traverses all leaf nodes.
+        // LoRA params are shared across groups, so we accumulate their gradients.
+        std::vector<at::Tensor> grad_inputs = {input};
+
+        if (!ctx->adapters.empty()) {
+            // Multi-LoRA: collect A/B from ctx->adapters
+            for (int64_t l = start; l < end; l++) {
+                int64_t lora_count = lora_pair_count(ctx->layer_configs[l]);
+                for (auto& adapter : ctx->adapters) {
+                    auto it = adapter.params.find(l);
+                    if (it == adapter.params.end()) continue;
+                    for (int64_t k = 0; k < lora_count && k < (int64_t)it->second.size(); k++) {
+                        auto& [a, b] = it->second[k];
+                        if (!a.requires_grad() && !b.requires_grad()) continue;
+                        grad_inputs.push_back(a);
+                        grad_inputs.push_back(b);
+                    }
+                }
+            }
+        } else {
+            // Legacy single-LoRA
+            for (int64_t l = start; l < end; l++) {
+                int64_t lora_count = lora_pair_count(ctx->layer_configs[l]);
+                int64_t la_offset = ctx->lora_layer_offset[l];
+                bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size();
+                if (has_lora) {
+                    for (int64_t k = 0; k < lora_count; k++) {
+                        if (!legacy_lora_slot_active(ctx, la_offset + k)) continue;
+                        grad_inputs.push_back(ctx->lora_a[la_offset + k]);
+                        grad_inputs.push_back(ctx->lora_b[la_offset + k]);
+                    }
+                }
+            }
+        }
+
+        auto grads = torch::autograd::grad(
+            {output}, grad_inputs, {grad},
+            /*retain_graph=*/false, /*create_graph=*/false,
+            /*allow_unused=*/true
+        );
+
+
+        // Manually accumulate LoRA param gradients
+        if (!ctx->adapters.empty()) {
+            // Multi-LoRA: accumulate into ctx->adapters
+            int64_t gi = 1;  // skip input grad (index 0)
+            for (int64_t l = start; l < end; l++) {
+                int64_t lora_count = lora_pair_count(ctx->layer_configs[l]);
+                for (auto& adapter : ctx->adapters) {
+                    auto it = adapter.params.find(l);
+                    if (it == adapter.params.end()) continue;
+                    for (int64_t k = 0; k < lora_count && k < (int64_t)it->second.size(); k++) {
+                        auto& [a, b] = it->second[k];
+                        if (!a.requires_grad() && !b.requires_grad()) continue;
+                        if (gi < (int64_t)grads.size() && grads[gi].defined()) {
+                            if (a.grad().defined()) a.grad().add_(grads[gi]);
+                            else a.mutable_grad() = grads[gi].clone();
+                        }
+                        gi++;
+                        if (gi < (int64_t)grads.size() && grads[gi].defined()) {
+                            if (b.grad().defined()) b.grad().add_(grads[gi]);
+                            else b.mutable_grad() = grads[gi].clone();
+                        }
+                        gi++;
+                    }
+                }
+            }
+        } else {
+            // Legacy single-LoRA
+            int64_t gi = 1;  // skip input grad (index 0)
+            for (int64_t l = start; l < end; l++) {
+                int64_t lora_count = lora_pair_count(ctx->layer_configs[l]);
+                int64_t la_offset = ctx->lora_layer_offset[l];
+                bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size();
+                if (has_lora) {
+                    for (int64_t k = 0; k < lora_count; k++) {
+                        if (!legacy_lora_slot_active(ctx, la_offset + k)) continue;
+                        if (grads[gi].defined()) {
+                            auto& pa = ctx->lora_a[la_offset + k];
+                            if (pa.grad().defined()) pa.grad().add_(grads[gi]);
+                            else pa.mutable_grad() = grads[gi].clone();
+                        }
+                        gi++;
+                        if (grads[gi].defined()) {
+                            auto& pb = ctx->lora_b[la_offset + k];
+                            if (pb.grad().defined()) pb.grad().add_(grads[gi]);
+                            else pb.mutable_grad() = grads[gi].clone();
+                        }
+                        gi++;
+                    }
+                }
+            }
+        }
+
+        // Gradient for this group's input = gradient for next group's output
+        grad = grads[0];
+
+        // Free saved input to release memory
+        ctx->group_inputs[g] = at::Tensor();
+    }
+}
+
+// ──────────────────────────────────────────────────────────────────────
+// Fused Cross-Entropy Loss with online softmax (FlashAttention-style).
+//
+// Instead of materializing [n_tokens, vocab] logits (8+ GB), we tile over
+// the vocabulary dimension:
+//   1. Forward: iterate vocab tiles, compute partial logits, accumulate
+//      global max + sum_exp via online softmax → loss
+//   2. Backward: iterate vocab tiles again, compute softmax = exp(logit-max)/sum_exp,
+//      subtract one_hot for target tokens → accumulate grad into hidden_normed
+//
+// Peak memory: [n_tokens, tile_size] instead of [n_tokens, vocab].
+// For N=100, seq=512: n_tokens=51200, vocab=248320, tile=8192
+//   Old: [51200, 248320] × 4 bytes = 49 GB per chunk
+//   New: [51200, 8192]  × 4 bytes = 1.6 GB per tile
+//
+// Returns: scalar loss tensor (with autograd graph for hidden_normed)
+// ──────────────────────────────────────────────────────────────────────
+struct LossResult {
+    at::Tensor value;
+    at::Tensor hidden_grad;
+    at::Tensor sample_loss_numerators;
+    at::Tensor sample_token_counts;
+};
+
+// Vocabulary-parallel cross entropy. The first LM-head pass computes local
+// online-softmax statistics and target logits. TP MAX/SUM reductions produce
+// global statistics; the backward projection reuses cached local logits when
+// they fit under a bounded workspace cap, otherwise it recomputes them.
+static LossResult compute_vocab_parallel_loss(
+    TrainingContext* ctx,
+    const at::Tensor& hidden,
+    const at::Tensor& input_ids,
+    const at::Tensor& target_mask,
+    bool independent_samples,
+    bool compute_hidden_grad,
+    bool collect_sample_losses = false
+) {
+    TORCH_CHECK(vocab_parallel_enabled(ctx),
+        "distributed vocabulary loss requires vocabulary TP");
+    TORCH_CHECK(ctx->tp_comm, "distributed vocabulary loss requires a TP communicator");
+
+    auto final_norm = *ctx->final_norm_ptr[0];
+    auto lm_head = *ctx->lm_head_ptr[0];
+    auto hidden_detached = hidden.detach();
+    if (compute_hidden_grad) hidden_detached.set_requires_grad(true);
+
+    at::Tensor hidden_normed;
+    {
+        at::AutoGradMode no_grad(false);
+        hidden_normed = rms_norm(hidden_detached, final_norm, ctx->rms_eps);
+    }
+
+    const int64_t batch_size = hidden_normed.size(0);
+    const int64_t seq_len = hidden_normed.size(1);
+    const int64_t hidden_dim = hidden_normed.size(2);
+    const int64_t shifted_seq_len = seq_len - 1;
+    auto hidden_flat = hidden_normed.narrow(1, 0, shifted_seq_len)
+        .reshape({-1, hidden_dim}).contiguous();
+    auto shifted_targets = input_ids.narrow(1, 1, shifted_seq_len).reshape({-1});
+    auto shifted_mask = target_mask.narrow(1, 1, shifted_seq_len)
+        .reshape({-1}).to(at::kFloat);
+    const int64_t total_tokens = shifted_targets.size(0);
+
+    int64_t token_tile = 512;
+    if (const char* value = getenv("QWEN36_CE_TOKEN_TILE")) {
+        token_tile = std::max<int64_t>(1, std::atoll(value));
+    }
+    int64_t vocab_tile = ctx->local_vocab_size;
+    if (const char* value = getenv("QWEN36_CE_TILE")) {
+        vocab_tile = std::max<int64_t>(1, std::atoll(value));
+    }
+    int64_t logits_cache_bytes = 512LL * 1024LL * 1024LL;
+    if (const char* value = getenv("QWEN36_CE_LOGITS_CACHE_BYTES")) {
+        logits_cache_bytes = std::max<int64_t>(0, std::atoll(value));
+    }
+
+    auto total_count = shifted_mask.sum().clamp_min(1.0);
+    at::Tensor token_denominators;
+    at::Tensor sample_loss_numerators;
+    at::Tensor sample_token_counts;
+    if (independent_samples) {
+        auto per_sample_count = target_mask.narrow(1, 1, shifted_seq_len)
+            .sum(1).to(at::kFloat);
+        token_denominators = per_sample_count.clamp_min(1.0).reshape({batch_size, 1})
+            .expand({batch_size, shifted_seq_len}).reshape({-1}).to(at::kFloat);
+        if (collect_sample_losses) {
+            sample_token_counts = per_sample_count;
+            sample_loss_numerators = at::zeros({batch_size},
+                at::TensorOptions().dtype(at::kFloat).device(hidden.device()));
+        }
+    }
+
+    auto total_loss = at::zeros({1},
+        at::TensorOptions().dtype(at::kFloat).device(hidden.device()));
+    at::Tensor grad_hidden_flat;
+    if (compute_hidden_grad) {
+        grad_hidden_flat = at::empty({total_tokens, hidden_dim},
+            at::TensorOptions().dtype(at::kFloat).device(hidden.device()));
+    }
+
+    at::AutoGradMode no_grad(false);
+    const int64_t vocab_start = ctx->tp_rank * ctx->local_vocab_size;
+    const int64_t num_token_tiles =
+        (total_tokens + token_tile - 1) / token_tile;
+    const int64_t num_vocab_tiles =
+        (ctx->local_vocab_size + vocab_tile - 1) / vocab_tile;
+
+    for (int64_t token_index = 0; token_index < num_token_tiles; ++token_index) {
+        const int64_t token_start = token_index * token_tile;
+        const int64_t token_count = std::min(
+            token_tile, total_tokens - token_start);
+        auto chunk_hidden = hidden_flat.narrow(0, token_start, token_count);
+        auto chunk_targets = shifted_targets.narrow(0, token_start, token_count);
+        auto chunk_mask = shifted_mask.narrow(0, token_start, token_count);
+
+        auto local_max = at::full({token_count, 1},
+            -std::numeric_limits<float>::infinity(),
+            at::TensorOptions().dtype(at::kFloat).device(hidden.device()));
+        auto local_sum_exp = at::zeros_like(local_max);
+        auto local_target_logit = at::zeros_like(local_max);
+        const bool cache_logits =
+            token_count * ctx->local_vocab_size * static_cast<int64_t>(sizeof(float)) <=
+            logits_cache_bytes;
+        std::vector<at::Tensor> cached_logits;
+        if (cache_logits) cached_logits.reserve(num_vocab_tiles);
+
+        for (int64_t vocab_index = 0; vocab_index < num_vocab_tiles; ++vocab_index) {
+            const int64_t local_start = vocab_index * vocab_tile;
+            const int64_t local_count = std::min(
+                vocab_tile, ctx->local_vocab_size - local_start);
+            const int64_t global_start = vocab_start + local_start;
+            const int64_t global_end = global_start + local_count;
+            auto head_tile = lm_head.narrow(0, local_start, local_count);
+            auto logits = at::matmul(chunk_hidden, head_tile.t()).to(at::kFloat);
+            if (cache_logits) cached_logits.push_back(logits);
+
+            auto tile_max = std::get<0>(at::max(logits, 1, true));
+            auto new_max = at::max(local_max, tile_max);
+            local_sum_exp = at::exp(local_max - new_max) * local_sum_exp +
+                at::exp(logits - new_max).sum(1, true);
+            local_max = new_max;
+
+            auto in_range = (chunk_targets >= global_start) &
+                (chunk_targets < global_end);
+            auto local_targets = (chunk_targets - global_start)
+                .clamp(0, local_count - 1).reshape({-1, 1});
+            auto gathered = at::gather(logits, 1, local_targets);
+            local_target_logit.add_(at::where(
+                in_range.reshape({-1, 1}), gathered, at::zeros_like(gathered)));
+        }
+
+        auto global_max = tp_allreduce_value(ctx, local_max, ncclMax);
+        local_sum_exp.mul_(at::exp(local_max - global_max));
+        auto global_stats = tp_allreduce_value(
+            ctx, at::cat({local_sum_exp, local_target_logit}, 1), ncclSum);
+        auto global_sum_exp = global_stats.narrow(1, 0, 1);
+        auto target_logit = global_stats.narrow(1, 1, 1);
+        auto per_token_loss =
+            (at::log(global_sum_exp) + global_max - target_logit).squeeze(1);
+        auto masked_loss = per_token_loss * chunk_mask;
+        if (independent_samples) {
+            if (sample_loss_numerators.defined()) {
+                auto token_indexes = at::arange(
+                    token_start, token_start + token_count,
+                    shifted_targets.options());
+                auto sample_indexes = at::floor_divide(
+                    token_indexes, shifted_seq_len);
+                sample_loss_numerators.index_add_(
+                    0, sample_indexes, masked_loss.detach());
+            }
+            total_loss.add_((masked_loss /
+                token_denominators.narrow(0, token_start, token_count)).sum());
+        } else {
+            total_loss.add_(masked_loss.sum() / total_count);
+        }
+
+        if (!compute_hidden_grad) continue;
+
+        auto grad_scale = independent_samples
+            ? chunk_mask /
+                token_denominators.narrow(0, token_start, token_count)
+            : chunk_mask / total_count;
+        auto local_grad_hidden = at::zeros({token_count, hidden_dim},
+            at::TensorOptions().dtype(at::kFloat).device(hidden.device()));
+        for (int64_t vocab_index = 0; vocab_index < num_vocab_tiles; ++vocab_index) {
+            const int64_t local_start = vocab_index * vocab_tile;
+            const int64_t local_count = std::min(
+                vocab_tile, ctx->local_vocab_size - local_start);
+            const int64_t global_start = vocab_start + local_start;
+            const int64_t global_end = global_start + local_count;
+            auto head_tile = lm_head.narrow(0, local_start, local_count);
+            auto logits = cache_logits
+                ? cached_logits[vocab_index]
+                : at::matmul(chunk_hidden, head_tile.t()).to(at::kFloat);
+            auto grad_logits = at::exp(logits - global_max) / global_sum_exp;
+
+            auto in_range = (chunk_targets >= global_start) &
+                (chunk_targets < global_end);
+            auto local_targets = (chunk_targets - global_start)
+                .clamp(0, local_count - 1).reshape({-1, 1});
+            auto one_hot = at::zeros_like(grad_logits);
+            one_hot.scatter_(1, local_targets, 1.0);
+            grad_logits.sub_(one_hot * in_range.to(at::kFloat).reshape({-1, 1}));
+            grad_logits.mul_(grad_scale.reshape({-1, 1}));
+            local_grad_hidden.add_(at::matmul(
+                grad_logits.to(head_tile.scalar_type()), head_tile).to(at::kFloat));
+        }
+        auto global_grad_hidden = tp_allreduce_value(
+            ctx, local_grad_hidden, ncclSum);
+        grad_hidden_flat.narrow(0, token_start, token_count)
+            .copy_(global_grad_hidden);
+    }
+
+    at::Tensor hidden_grad;
+    if (compute_hidden_grad) {
+        auto grad_shifted = grad_hidden_flat
+            .reshape({batch_size, shifted_seq_len, hidden_dim})
+            .to(hidden_normed.scalar_type());
+        auto grad_full = at::cat({
+            grad_shifted,
+            at::zeros({batch_size, 1, hidden_dim},
+                at::TensorOptions().dtype(hidden_normed.scalar_type())
+                    .device(hidden_normed.device()))
+        }, 1);
+        at::AutoGradMode grad_mode(true);
+        auto hidden_normed_recompute = rms_norm(
+            hidden_detached, final_norm, ctx->rms_eps);
+        hidden_grad = torch::autograd::grad(
+            {hidden_normed_recompute}, {hidden_detached}, {grad_full},
+            /*retain_graph=*/false, /*create_graph=*/false,
+            /*allow_unused=*/false)[0];
+    }
+
+    return {
+        total_loss,
+        hidden_grad,
+        sample_loss_numerators,
+        sample_token_counts,
+    };
+}
+
+static at::Tensor compute_loss_fused(
+    TrainingContext* ctx,
+    const at::Tensor& hidden,       // [batch, seq, hidden] (requires_grad)
+    const at::Tensor& input_ids,    // [batch, seq]
+    const at::Tensor& target_mask,  // [batch, seq]
+    int64_t vocab_size
+) {
+    if (vocab_parallel_enabled(ctx)) {
+        return compute_vocab_parallel_loss(
+            ctx, hidden, input_ids, target_mask,
+            /*independent_samples=*/false,
+            /*compute_hidden_grad=*/false).value;
+    }
+    auto final_norm = *ctx->final_norm_ptr[0];
+    auto lm_head = *ctx->lm_head_ptr[0];  // [vocab, hidden]
+
+    // Compute hidden_normed in no-grad, then set requires_grad.
+    auto hidden_detached = hidden.detach();
+    at::Tensor hidden_normed;
+    {
+        at::AutoGradMode no_grad_mode(false);
+        hidden_normed = rms_norm(hidden_detached, final_norm, ctx->rms_eps);
+    }
+    hidden_normed.set_requires_grad(true);
+
+    int64_t seq_len = hidden_normed.size(1);
+    int64_t hidden_dim = hidden_normed.size(2);
+
+    auto shifted_hidden = hidden_normed.narrow(1, 0, seq_len - 1);
+    auto shifted_targets = input_ids.narrow(1, 1, seq_len - 1).reshape({-1});
+    auto shifted_mask = target_mask.narrow(1, 1, seq_len - 1).reshape({-1});
+
+    int64_t total_tokens = shifted_targets.size(0);
+    auto hidden_flat = shifted_hidden.reshape({-1, hidden_dim});
+
+    // Ensure hidden_flat is contiguous for narrow + matmul
+    hidden_flat = hidden_flat.contiguous();
+
+    auto mask_f = shifted_mask.to(at::kFloat);
+
+    // ── Tile configuration ──
+    // tile_size controls vocab granularity. Larger = fewer iterations but more memory.
+    // [total_tokens, tile] × 4 bytes (FP32). 8192 → ~1.6 GB for 51200 tokens.
+    int64_t tile_size = 8192;
+    const char* ts_env = getenv("QWEN36_CE_TILE");
+    if (ts_env) { tile_size = atol(ts_env); if (tile_size < 1) tile_size = 8192; }
+    int64_t num_tiles = (vocab_size + tile_size - 1) / tile_size;
+
+    // lm_head transpose: we need [hidden, vocab] for matmul
+    // lm_head is [vocab, hidden], so lm_head.t() is [hidden, vocab]
+    // We narrow on dim 0 of lm_head (vocab dim), then transpose.
+    // lm_head_w: [tile, hidden] → .t() → [hidden, tile]
+    // hidden_flat: [total_tokens, hidden] × [hidden, tile] → [total_tokens, tile]
+
+    // ── Forward pass: compute loss via online softmax ──
+    // For each token i: loss_i = log(sum_v exp(logit_iv)) - logit_i,target_i
+    // We compute in two phases:
+    //   Phase 1: find global max and sum_exp across all vocab tiles
+    //   Phase 2: compute loss = log(sum_exp) - target_logit (for masked tokens)
+
+    // Running max and sum_exp per token
+    auto logit_max = at::full({total_tokens, 1}, -std::numeric_limits<float>::infinity(),
+        at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device()));
+    auto sum_exp = at::zeros({total_tokens, 1},
+        at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device()));
+
+    // Phase 1: accumulate max and sum_exp
+    for (int64_t t = 0; t < num_tiles; t++) {
+        int64_t v_start = t * tile_size;
+        int64_t v_end = std::min(v_start + tile_size, vocab_size);
+        int64_t v_n = v_end - v_start;
+
+        auto lm_head_tile = lm_head.narrow(0, v_start, v_n);  // [v_n, hidden]
+        // [total_tokens, hidden] × [hidden, v_n] → [total_tokens, v_n]
+        auto logits_tile = at::matmul(hidden_flat, lm_head_tile.t()).to(at::kFloat);
+
+        // Online softmax update: new_max = max(old_max, tile_max)
+        auto tile_max = std::get<0>(at::max(logits_tile, /*dim=*/1, /*keepdim=*/true));
+        auto new_max = at::max(logit_max, tile_max);
+
+        // Adjust sum_exp: exp(old - new_max) * old_sum + exp(tile - new_max) * tile_sum
+        auto old_exp = at::exp(logit_max - new_max);
+        auto tile_exp = at::exp(logits_tile - new_max);
+        sum_exp = old_exp * sum_exp + tile_exp.sum(/*dim=*/1, /*keepdim=*/true);
+        logit_max = new_max;
+
+    }
+
+    // Phase 2: compute target logit for each token
+    // We need logits[target] — gather the target position's logit.
+    // Do a second pass to find target logit.
+    auto target_logit = at::full({total_tokens, 1},
+        -std::numeric_limits<float>::infinity(),
+        at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device()));
+
+    for (int64_t t = 0; t < num_tiles; t++) {
+        int64_t v_start = t * tile_size;
+        int64_t v_end = std::min(v_start + tile_size, vocab_size);
+        int64_t v_n = v_end - v_start;
+
+        // Check if any target falls in this tile's vocab range
+        auto in_range = (shifted_targets >= v_start) & (shifted_targets < v_end);
+        if (!in_range.any().item<bool>()) continue;
+
+        auto lm_head_tile = lm_head.narrow(0, v_start, v_n);
+        auto logits_tile = at::matmul(hidden_flat, lm_head_tile.t()).to(at::kFloat);
+
+        // Gather target logits: subtract v_start to get local index
+        auto local_targets = (shifted_targets - v_start).clamp(0, v_n - 1);
+        // Gather: logits_tile[i, local_targets[i]] for tokens in range
+        auto gathered = at::gather(logits_tile, /*dim=*/1, local_targets.reshape({-1, 1}));
+        // Only keep tokens that are actually in range
+        gathered = at::where(
+            in_range.reshape({-1, 1}), gathered,
+            at::full_like(gathered, -std::numeric_limits<float>::infinity()));
+        target_logit = at::max(target_logit, gathered);
+
+    }
+
+    // Loss per token: log(sum_exp) - target_logit  (= -log(softmax[target]))
+    auto log_sum_exp = at::log(sum_exp) + logit_max;  // log(sum(exp(x-max))) + max = logsumexp
+    auto per_token_loss = log_sum_exp - target_logit;  // [total_tokens, 1]
+    per_token_loss = at::where(
+        mask_f.reshape({-1, 1}) > 0, per_token_loss,
+        at::zeros_like(per_token_loss));
+    auto masked_loss = per_token_loss.squeeze(1) * mask_f;
+    auto total_count = mask_f.sum().clamp_min(1.0);
+    double loss_val = (masked_loss.sum().item<double>()) / total_count.item<double>();
+
+    // Evaluation runs with AutoGradMode disabled. The online-softmax passes
+    // above already produced the exact scalar value, so do not execute the
+    // training-only manual gradient pass below. Besides fixing eval on a
+    // no-grad graph, this avoids a third traversal of the vocabulary tiles.
+    if (!at::GradMode::is_enabled()) {
+        return at::tensor({loss_val},
+            at::TensorOptions().dtype(at::kFloat).device(hidden.device()));
+    }
+
+    // ── Backward pass: compute grad_hidden_normed manually ──
+    // dL/dhidden_normed = (softmax - one_hot) / count * mask
+    // softmax = exp(logit - logit_max) / sum_exp
+    // We iterate tiles again, accumulate grad = softmax_tile @ lm_head_tile + target_grad
+    auto grad_hidden = at::zeros({total_tokens, hidden_dim},
+        at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device()));
+    auto grad_scale = mask_f / total_count;  // [total_tokens]
+
     for (int64_t t = 0; t < num_tiles; t++) {
         int64_t v_start = t * tile_size;
         int64_t v_end = std::min(v_start + tile_size, vocab_size);
         int64_t v_n = v_end - v_start;
 
-        // Check if any target falls in this tile's vocab range
-        auto in_range = (shifted_targets >= v_start) & (shifted_targets < v_end);
-        if (!in_range.any().item<bool>()) continue;
+        auto lm_head_tile = lm_head.narrow(0, v_start, v_n);  // [v_n, hidden]
+        auto logits_tile = at::matmul(hidden_flat, lm_head_tile.t()).to(at::kFloat);
+
+        // softmax_tile = exp(logits - max) / sum_exp  (reuse from forward)
+        auto softmax_tile = at::exp(logits_tile - logit_max) / sum_exp;  // [total_tokens, v_n]
+
+        // Subtract one_hot for target tokens in this tile
+        auto in_range = (shifted_targets >= v_start) & (shifted_targets < v_end);
+        if (in_range.any().item<bool>()) {
+            auto local_targets = (shifted_targets - v_start).clamp_min(0);
+            // scatter -1 at target positions
+            auto one_hot = at::zeros({total_tokens, v_n},
+                at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device()));
+            one_hot.scatter_(1, local_targets.reshape({-1, 1}), 1.0);
+            one_hot = one_hot * in_range.to(at::kFloat).reshape({-1, 1});
+            softmax_tile = softmax_tile - one_hot;
+        }
+
+        // grad_hidden += grad_scale * softmax_tile @ lm_head_tile
+        // softmax_tile: [total_tokens, v_n] (Float), lm_head_tile: [v_n, hidden] (BF16)
+        // → [total_tokens, hidden] (Float)
+        auto grad_tile = at::matmul(
+            (softmax_tile * grad_scale.reshape({-1, 1})),
+            lm_head_tile.to(at::kFloat)
+        );
+        grad_hidden.add_(grad_tile);
+
+    }
+
+    // Set gradient on hidden_normed (leaf tensor).
+    // grad_hidden covers [batch, seq-1, hidden] (shifted tokens).
+    // hidden_normed is [batch, seq, hidden] — the final position has no target.
+    auto grad_reshaped = grad_hidden.to(hidden_normed.scalar_type())
+        .reshape({hidden_normed.size(0), seq_len - 1, hidden_dim});
+    auto grad_full = at::cat({
+        grad_reshaped,
+        at::zeros({hidden_normed.size(0), 1, hidden_dim},
+            at::TensorOptions().dtype(hidden_normed.scalar_type()).device(hidden_normed.device()))
+    }, /*dim=*/1);  // [batch, seq, hidden]
+    hidden_normed.mutable_grad() = grad_full;
+
+    // Backprop hidden_normed gradient to hidden via rms_norm recompute.
+    if (hidden_normed.grad().defined()) {
+        hidden.set_requires_grad(true);
+        auto hidden_normed_recompute = rms_norm(hidden, final_norm, ctx->rms_eps);
+        hidden_normed_recompute.backward(hidden_normed.grad());
+    }
+
+    return at::tensor({loss_val},
+        at::TensorOptions().dtype(at::kFloat).device(hidden.device()));
+}
+
+// Cross-entropy loss with response-only masking — chunked with detach.
+// Return the hidden gradient instead of mutating/backpropagating through the
+// main model graph. The caller combines auxiliary gradients and performs one
+// main backward, which is required for checkpointed execution.
+static LossResult compute_loss(
+    TrainingContext* ctx,
+    const at::Tensor& hidden,
+    const at::Tensor& input_ids,
+    const at::Tensor& target_mask,
+    int64_t vocab_size,
+    bool independent_samples = false,
+    bool collect_sample_losses = false
+) {
+    if (vocab_parallel_enabled(ctx)) {
+        return compute_vocab_parallel_loss(
+            ctx, hidden, input_ids, target_mask,
+            independent_samples, /*compute_hidden_grad=*/true,
+            collect_sample_losses);
+    }
+    auto final_norm = *ctx->final_norm_ptr[0];
+    auto lm_head = *ctx->lm_head_ptr[0];
+
+    // Detach hidden for CE computation — CE backward won't touch main model graph.
+    // We accumulate gradient into hidden_normed, then manually backprop to hidden.
+    auto hidden_detached = hidden.detach().set_requires_grad(true);
+
+    // Compute hidden_normed in no-grad, then set requires_grad on it.
+    // This way CE backward only builds a tiny graph (hidden_normed → logits → loss),
+    // not connected to hidden_detached at all.
+    at::Tensor hidden_normed;
+    {
+        at::AutoGradMode no_grad_mode(false);
+        hidden_normed = rms_norm(hidden_detached, final_norm, ctx->rms_eps);
+    }
+    hidden_normed.set_requires_grad(true);
+
+    int64_t seq_len = hidden_normed.size(1);
+    auto shifted_hidden = hidden_normed.narrow(1, 0, seq_len - 1);
+    auto shifted_targets = input_ids.narrow(1, 1, seq_len - 1).reshape({-1});
+    auto shifted_mask = target_mask.narrow(1, 1, seq_len - 1).reshape({-1});
+
+    int64_t total_tokens = shifted_targets.size(0);
+    // Smaller chunks = less peak memory (4GB vs 16GB per chunk at vocab=248K).
+    // This allows removing emptyCache between chunks — GPU stays async.
+    int64_t chunk_size = 4096;
+    int64_t num_chunks = (total_tokens + chunk_size - 1) / chunk_size;
+
+    auto total_count = shifted_mask.sum().clamp_min(1.0);
+    at::Tensor token_denominators;
+    at::Tensor sample_loss_numerators;
+    at::Tensor sample_token_counts;
+    if (independent_samples) {
+        auto per_sample_count = target_mask.narrow(1, 1, seq_len - 1)
+            .sum(1).to(at::kFloat);
+        token_denominators = per_sample_count.clamp_min(1.0)
+            .reshape({target_mask.size(0), 1})
+            .expand({target_mask.size(0), seq_len - 1}).reshape({-1});
+        if (collect_sample_losses) {
+            sample_token_counts = per_sample_count;
+            sample_loss_numerators = at::zeros({target_mask.size(0)},
+                at::TensorOptions().dtype(at::kFloat).device(hidden.device()));
+        }
+    }
+    auto hidden_flat = shifted_hidden.reshape({-1, hidden_normed.size(2)});
+
+    double total_loss_val = 0.0;
+
+    for (int64_t c = 0; c < num_chunks; c++) {
+        int64_t start = c * chunk_size;
+        int64_t end = std::min(start + chunk_size, total_tokens);
+        int64_t n = end - start;
+
+        auto chunk_hidden = hidden_flat.narrow(0, start, n);
+        auto chunk_logits = at::matmul(chunk_hidden, lm_head.t());
+        auto chunk_targets = shifted_targets.narrow(0, start, n);
+        auto chunk_mask = shifted_mask.narrow(0, start, n);
+
+        // Diagnostic: print logits stats for first chunk, first token
+        if (c == 0 && getenv("QWEN36_LOSS_DIAG")) {
+            auto logits_f = chunk_logits[0].to(at::kFloat);
+            auto lsm = at::log_softmax(logits_f, -1);
+            int64_t tgt = chunk_targets[0].item<int64_t>();
+            fprintf(stderr, "[diag] logits shape: [%ld, %ld]\n", (long)n, (long)chunk_logits.size(1));
+            fprintf(stderr, "[diag] logits[0,:5]: %.6f %.6f %.6f %.6f %.6f\n",
+                    logits_f[0].item<float>(), logits_f[1].item<float>(),
+                    logits_f[2].item<float>(), logits_f[3].item<float>(),
+                    logits_f[4].item<float>());
+            fprintf(stderr, "[diag] logits mean: %.6f, std: %.6f\n",
+                    logits_f.mean().item<float>(), logits_f.std().item<float>());
+            fprintf(stderr, "[diag] target token: %ld\n", (long)tgt);
+            fprintf(stderr, "[diag] target logit: %.6f\n", logits_f[tgt].item<float>());
+            fprintf(stderr, "[diag] log_softmax[target]: %.6f\n", lsm[tgt].item<float>());
+            fprintf(stderr, "[diag] -log_softmax[target] (per-token loss): %.6f\n", -lsm[tgt].item<float>());
+        }
+
+        auto per_token_loss = at::cross_entropy_loss(
+            chunk_logits.to(at::kFloat), chunk_targets,
+            at::Tensor(), at::Reduction::None, -100, 0.0
+        );
+        auto masked_loss = per_token_loss * chunk_mask.to(at::kFloat);
+        if (sample_loss_numerators.defined()) {
+            auto token_indexes = at::arange(
+                start, end, shifted_targets.options());
+            auto sample_indexes = at::floor_divide(token_indexes, seq_len - 1);
+            sample_loss_numerators.index_add_(
+                0, sample_indexes, masked_loss.detach());
+        }
+        // Single-sample training uses the global response-token mean. Batched
+        // multi-LoRA instead divides each row by its own response-token count
+        // so tenant gradients do not depend on neighboring rows.
+        auto chunk_loss = independent_samples
+            ? (masked_loss / token_denominators.narrow(0, start, n)).sum()
+            : masked_loss.sum() / total_count;
+
+        // Backward this chunk — each chunk creates an independent CE subgraph
+        // because hidden_normed is a leaf tensor. retain_graph=false is safe
+        // and much faster than retain_graph=true (which accumulates graph).
+        torch::autograd::backward({chunk_loss}, {},
+            /*retain_graph=*/false, /*create_graph=*/false);
+
+        total_loss_val += chunk_loss.item<double>();
+
+    }
+
+    TORCH_CHECK(hidden_normed.grad().defined(),
+        "cross-entropy did not produce a hidden gradient");
+    auto hidden_normed_recompute = rms_norm(hidden_detached, final_norm, ctx->rms_eps);
+    auto hidden_grad = torch::autograd::grad(
+        {hidden_normed_recompute}, {hidden_detached}, {hidden_normed.grad()},
+        /*retain_graph=*/false, /*create_graph=*/false,
+        /*allow_unused=*/false)[0];
+
+    return {
+        at::tensor({total_loss_val},
+            at::TensorOptions().dtype(at::kFloat).device(hidden.device())),
+        hidden_grad,
+        sample_loss_numerators,
+        sample_token_counts,
+    };
+}
+
+static PipelineLossResult pipeline_compute_loss(
+    TrainingContext* ctx,
+    const at::Tensor& hidden,
+    const at::Tensor& input_ids,
+    const at::Tensor& target_mask
+) {
+    auto result = compute_loss(
+        ctx, hidden, input_ids, target_mask, ctx->vocab_size,
+        /*independent_samples=*/ctx->pp_window.dynamic_lora,
+        /*collect_sample_losses=*/ctx->pp_window.dynamic_lora);
+    double value = result.value.item<double>();
+    double numerator = value;
+    double denominator = 1.0;
+    if (ctx->pp_window.dynamic_lora && result.sample_loss_numerators.defined() &&
+            result.sample_token_counts.defined()) {
+        numerator = result.sample_loss_numerators.sum().item<double>();
+        denominator = result.sample_token_counts.sum().item<double>();
+        value = numerator / std::max(1.0, denominator);
+    }
+    return {
+        value,
+        result.hidden_grad,
+        numerator,
+        denominator,
+        result.sample_loss_numerators,
+        result.sample_token_counts,
+    };
+}
+
+// ──────────────────────────────────────────────────────────────────────
+// MTP (Multi-Token Prediction) forward + loss
+// ──────────────────────────────────────────────────────────────────────
+
+// MTP forward: produce hidden states (not logits) for chunked loss computation.
+// hidden: [batch, seq, hidden] — pre-norm hidden from main model
+// Returns: [batch, seq-1, hidden] — MTP hidden (after final norm, before lm_head)
+static at::Tensor mtp_forward(
+    TrainingContext* ctx,
+    const at::Tensor& hidden,
+    const at::Tensor& input_ids
+) {
+    auto kind = ctx->compute_type;
+    int64_t seq_len = hidden.size(1);
+
+    // hidden[t] + embed[t+1] → predict token t+2 (Megatron convention)
+    auto hidden_shifted = hidden.narrow(1, 0, seq_len - 1);  // [batch, seq-1, hidden]
+    auto embed_next = vocabulary_embedding(
+        ctx, input_ids.narrow(1, 1, seq_len - 1));  // [batch, seq-1, hidden]
+
+    // RMSNorm both
+    auto h_normed = rms_norm(hidden_shifted, *ctx->mtp_pre_fc_norm_hidden, ctx->rms_eps).to(kind);
+    auto e_normed = rms_norm(embed_next, *ctx->mtp_pre_fc_norm_emb, ctx->rms_eps).to(kind);
+
+    // Combine: embed first, then hidden → fc projection
+    auto combined = at::cat({e_normed, h_normed}, /*dim=*/-1);
+    auto projected = at::matmul(combined, ctx->mtp_fc->t());  // fc: [hidden, 2*hidden]
+
+    // MTP layers (full attention + MoE/dense, no LoRA)
+    at::Tensor h = projected;
+    int64_t num_mtp_layers = (int64_t)ctx->mtp_layer_configs.size();
+    for (int64_t i = 0; i < num_mtp_layers; i++) {
+        int64_t w_offset = 0;
+        for (int64_t j = 0; j < i; j++)
+            w_offset += weight_count_for_layer(ctx->mtp_layer_configs[j]);
+        int64_t w_count = weight_count_for_layer(ctx->mtp_layer_configs[i]);
+        std::vector<at::Tensor*> layer_w(ctx->mtp_layer_weights.begin() + w_offset,
+                                         ctx->mtp_layer_weights.begin() + w_offset + w_count);
+        // MTP processes seq-1 tokens — slice attention mask's last dim to match
+        auto mtp_mask = ctx->attention_mask.defined()
+            ? ctx->attention_mask.narrow(-1, 0, h.size(1))
+            : at::Tensor();
+        auto mtp_lengths = ctx->attention_lengths.defined()
+            ? ctx->attention_lengths.clamp(0, h.size(1)).to(at::kInt).contiguous()
+            : at::Tensor();
+        h = forward_single_layer(ctx, h, layer_w.data(), &ctx->mtp_layer_configs[i],
+            ctx->num_layers + i, kind, mtp_mask, mtp_lengths,
+            ctx->base_tp_attention || ctx->base_tp_mlp);
+    }
+
+    // Final norm only — return hidden, not logits
+    return rms_norm(h, *ctx->mtp_norm, ctx->rms_eps).to(kind);
+}
+
+struct MtpLossResult {
+    at::Tensor value;
+    at::Tensor sample_loss_numerators;
+    at::Tensor sample_token_counts;
+};
+
+// Vocabulary-parallel MTP CE.  The online softmax statistics are computed
+// from detached local logits and reduced over TP; a zero-valued straight-
+// through surrogate supplies the exact local softmax gradient without asking
+// the generic NCCL helper to differentiate a collective it does not own.
+static MtpLossResult mtp_compute_vocab_parallel_loss(
+    TrainingContext* ctx,
+    const at::Tensor& mtp_hidden,
+    const at::Tensor& input_ids,
+    const at::Tensor& target_mask,
+    bool independent_samples,
+    bool collect_sample_losses
+) {
+    TORCH_CHECK(vocab_parallel_enabled(ctx) && ctx->tp_comm,
+        "vocabulary-parallel MTP requires an initialized TP communicator");
+    const int64_t seq_len = input_ids.size(1);
+    const int64_t n_tokens = seq_len - 2;
+    const int64_t batch_size = input_ids.size(0);
+    const int64_t hidden_dim = mtp_hidden.size(2);
+    auto hidden_flat = mtp_hidden.narrow(1, 0, n_tokens)
+        .reshape({-1, hidden_dim});
+    auto shifted_targets = input_ids.narrow(1, 2, n_tokens).reshape({-1});
+    auto shifted_mask = target_mask.narrow(1, 2, n_tokens)
+        .reshape({-1}).to(at::kFloat);
+    const int64_t total_tokens = shifted_targets.size(0);
+    const int64_t token_tile = 512;
+    const int64_t vocab_start = ctx->tp_rank * ctx->local_vocab_size;
+    const int64_t vocab_end = vocab_start + ctx->local_vocab_size;
+    auto lm_head = *ctx->lm_head_ptr[0];
+    auto total_count = shifted_mask.sum().clamp_min(1.0);
+    at::Tensor token_denominators;
+    at::Tensor sample_loss_numerators;
+    at::Tensor sample_token_counts;
+    if (independent_samples) {
+        auto per_sample_count = target_mask.narrow(1, 2, n_tokens)
+            .sum(1).to(at::kFloat);
+        token_denominators = per_sample_count.clamp_min(1.0)
+            .reshape({batch_size, 1}).expand({batch_size, n_tokens})
+            .reshape({-1});
+        if (collect_sample_losses) {
+            sample_token_counts = per_sample_count;
+            sample_loss_numerators = at::zeros({batch_size},
+                at::TensorOptions().dtype(at::kFloat)
+                    .device(mtp_hidden.device()));
+        }
+    }
+    auto total_loss = at::zeros({1},
+        at::TensorOptions().dtype(at::kFloat).device(mtp_hidden.device()));
+    for (int64_t start = 0; start < total_tokens; start += token_tile) {
+        const int64_t n = std::min(token_tile, total_tokens - start);
+        auto chunk_hidden = hidden_flat.narrow(0, start, n);
+        auto chunk_targets = shifted_targets.narrow(0, start, n);
+        auto chunk_mask = shifted_mask.narrow(0, start, n);
+        auto logits = at::matmul(chunk_hidden, lm_head.t()).to(at::kFloat);
+        auto logits_detached = logits.detach();
+        auto local_max = std::get<0>(at::max(logits_detached, 1, true));
+        auto global_max = tp_allreduce_value(ctx, local_max, ncclMax);
+        auto local_exp = at::exp(logits_detached - global_max);
+        auto global_sum = tp_allreduce_value(
+            ctx, local_exp.sum(1, true), ncclSum);
+        auto in_range = (chunk_targets >= vocab_start) &
+            (chunk_targets < vocab_end);
+        auto local_targets = (chunk_targets - vocab_start)
+            .clamp(0, ctx->local_vocab_size - 1).reshape({-1, 1});
+        auto local_target = at::gather(
+            logits_detached, 1, local_targets);
+        local_target = at::where(
+            in_range.reshape({-1, 1}), local_target,
+            at::zeros_like(local_target));
+        auto global_target = tp_allreduce_value(
+            ctx, local_target, ncclSum);
+        auto global_logsum = at::log(global_sum) + global_max;
+        auto global_per_loss = (global_logsum - global_target).squeeze(1);
+
+        auto local_probability = at::exp(
+            logits_detached - global_logsum);
+        auto local_one_hot = at::zeros_like(logits_detached);
+        local_one_hot.scatter_(1, local_targets,
+            in_range.reshape({-1, 1}).to(at::kFloat));
+        auto surrogate = (logits *
+            (local_probability - local_one_hot)).sum(1);
+        auto connected_loss = global_per_loss + surrogate - surrogate.detach();
+        auto masked_loss = connected_loss * chunk_mask;
+        if (sample_loss_numerators.defined()) {
+            auto token_indexes = at::arange(
+                start, start + n, shifted_targets.options());
+            auto sample_indexes = at::floor_divide(token_indexes, n_tokens);
+            sample_loss_numerators.index_add_(
+                0, sample_indexes, masked_loss.detach());
+        }
+        total_loss = total_loss + (independent_samples
+            ? (masked_loss / token_denominators.narrow(0, start, n)).sum()
+            : masked_loss.sum() / total_count);
+    }
+    return {
+        total_loss * ctx->mtp_loss_scale,
+        sample_loss_numerators.defined()
+            ? sample_loss_numerators * ctx->mtp_loss_scale
+            : at::Tensor(),
+        sample_token_counts,
+    };
+}
+
+// MTP loss: chunked matmul + cross-entropy.
+// mtp_hidden[t] (from hidden[t] + embed[t+1]) predicts token t+2 (Megatron convention)
+// No full logits tensor — chunked matmul + fused CE
+static MtpLossResult mtp_compute_loss(
+    TrainingContext* ctx,
+    const at::Tensor& mtp_hidden,
+    const at::Tensor& input_ids,
+    const at::Tensor& target_mask,
+    bool independent_samples = false,
+    bool collect_sample_losses = false
+) {
+    if (vocab_parallel_enabled(ctx)) {
+        return mtp_compute_vocab_parallel_loss(
+            ctx, mtp_hidden, input_ids, target_mask,
+            independent_samples, collect_sample_losses);
+    }
+    int64_t vocab_size = ctx->vocab_size;
+    int64_t seq_len = input_ids.size(1);
+    auto lm_head = *ctx->lm_head_ptr[0];
+
+    // MTP hidden: [batch, seq-1, hidden], drop last → predict t+2
+    int64_t n_tokens = seq_len - 2;
+    auto hidden_flat = mtp_hidden.narrow(1, 0, n_tokens).reshape({-1, mtp_hidden.size(2)});
+    auto shifted_targets = input_ids.narrow(1, 2, n_tokens).reshape({-1});
+    auto shifted_mask = target_mask.narrow(1, 2, n_tokens).reshape({-1});
+
+    // Chunked matmul + cross-entropy
+    int64_t total_tokens = shifted_targets.size(0);
+    int64_t chunk_size = 4096;  // smaller chunks = less peak memory
+    int64_t num_chunks = (total_tokens + chunk_size - 1) / chunk_size;
+
+    auto total_loss = at::zeros({1}, at::TensorOptions().dtype(at::kFloat).device(mtp_hidden.device()));
+    auto total_count = shifted_mask.sum().clamp_min(1.0);
+    at::Tensor token_denominators;
+    at::Tensor sample_loss_numerators;
+    at::Tensor sample_token_counts;
+    if (independent_samples) {
+        auto per_sample_count = target_mask.narrow(1, 2, n_tokens)
+            .sum(1).to(at::kFloat);
+        token_denominators = per_sample_count.clamp_min(1.0).reshape({-1, 1})
+            .expand({target_mask.size(0), n_tokens}).reshape({-1});
+        if (collect_sample_losses) {
+            sample_token_counts = per_sample_count;
+            sample_loss_numerators = at::zeros(
+                {target_mask.size(0)},
+                at::TensorOptions().dtype(at::kFloat).device(mtp_hidden.device()));
+        }
+    }
+
+    for (int64_t c = 0; c < num_chunks; c++) {
+        int64_t start = c * chunk_size;
+        int64_t end = std::min(start + chunk_size, total_tokens);
+        int64_t n = end - start;
+
+        auto chunk_hidden = hidden_flat.narrow(0, start, n);
+        auto chunk_logits = at::matmul(chunk_hidden, lm_head.t());  // [n, vocab]
+        auto chunk_targets = shifted_targets.narrow(0, start, n);
+        auto chunk_mask = shifted_mask.narrow(0, start, n);
+
+        auto per_token_loss = at::cross_entropy_loss(
+            chunk_logits.to(at::kFloat), chunk_targets,
+            /*weight=*/at::Tensor(), /*reduction=*/at::Reduction::None,
+            /*ignore_index=*/-100, /*label_smoothing=*/0.0
+        );
+        auto masked_loss = per_token_loss * chunk_mask.to(at::kFloat);
+        if (sample_loss_numerators.defined()) {
+            auto token_indexes = at::arange(
+                start, end, shifted_targets.options());
+            auto sample_indexes = at::floor_divide(token_indexes, n_tokens);
+            sample_loss_numerators.index_add_(
+                0, sample_indexes, masked_loss.detach());
+        }
+        // Avoid in-place accumulation into a non-grad leaf: out-of-place add
+        // keeps the MTP loss connected to the frozen-head input graph.
+        total_loss = total_loss + (independent_samples
+            ? (masked_loss / token_denominators.narrow(0, start, n)).sum()
+            : masked_loss.sum());
+    }
+
+    if (!independent_samples) total_loss = total_loss / total_count;
+    return {
+        total_loss * ctx->mtp_loss_scale,
+        sample_loss_numerators.defined()
+            ? sample_loss_numerators * ctx->mtp_loss_scale
+            : at::Tensor(),
+        sample_token_counts,
+    };
+}
+
+// ──────────────────────────────────────────────────────────────────────
+// C FFI
+// ──────────────────────────────────────────────────────────────────────
+
+extern "C" {
+
+__attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() {
+    return 31;
+}
+
+// Benchmark/metrics helper. Reducing over every orthogonal axis propagates the
+// maximum over the complete process grid without requiring a world group.
+__attribute__((visibility("default"))) double qwen36_parallel_max_double(
+    void* ctx_ptr, double value
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "parallel max requires a valid training context");
+        TORCH_CHECK(!ctx->topology_invalid,
+            "parallel max rejected an invalid process topology");
+        c10::cuda::set_device(ctx->cuda_device);
+        cudaSetDevice(ctx->cuda_device);
+        auto maximum = at::full(
+            {1}, value,
+            at::TensorOptions().device(at::kCUDA, ctx->cuda_device)
+                .dtype(at::kDouble));
+        auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+        auto reduce_axis = [&](ncclComm_t communicator, int axis_size,
+                               const char* axis) {
+            TORCH_CHECK(axis_size == 1 || communicator,
+                axis, " benchmark reduction is missing its communicator");
+            if (axis_size == 1) return;
+            auto err = ncclAllReduce(
+                maximum.data_ptr<double>(), maximum.data_ptr<double>(), 1,
+                ncclDouble, ncclMax, communicator, stream);
+            TORCH_CHECK(err == ncclSuccess, axis,
+                " benchmark max all-reduce failed: ", ncclGetErrorString(err));
+        };
+        reduce_axis(ctx->tp_comm, ctx->tp_world_size, "TP");
+        reduce_axis(ctx->cp_comm, ctx->cp_world_size, "CP");
+        reduce_axis(ctx->nccl_comm, ctx->ep_world_size, "EP");
+        reduce_axis(ctx->dp_comm, ctx->dp_world_size, "DP");
+        reduce_axis(ctx->pp_comm, ctx->pp_world_size, "PP");
+        return maximum.to(at::kCPU).item<double>();
+    } catch (const std::exception& error) {
+        fprintf(stderr, "[q36] parallel max FAILED: %s\n", error.what());
+        return std::numeric_limits<double>::quiet_NaN();
+    }
+}
+
+static constexpr int32_t QWEN36_CONTEXT_BASE_TP_ATTENTION = 1 << 0;
+static constexpr int32_t QWEN36_CONTEXT_DATA_PARALLEL = 1 << 1;
+static constexpr int32_t QWEN36_CONTEXT_VOCAB_PARALLEL = 1 << 2;
+static constexpr int32_t QWEN36_CONTEXT_EXPERT_PARALLEL = 1 << 3;
+static constexpr int32_t QWEN36_CONTEXT_BASE_TP_MLP = 1 << 4;
+
+// Create training context — called once at startup
+// lora_rank: LoRA rank (from config)
+// target_layers: array of layer indices to apply LoRA (nullptr = all layers)
+// num_target_layers: length of target_layers array
+static void* qwen36_create_training_context_impl(
+    void** weight_ptrs, int64_t num_weight_ptrs,
+    void* embed_ptr, void* final_norm_ptr, void* lm_head_ptr,
+    void* layer_configs_ptr, int64_t num_layers,
+    int64_t global_layer_start, int64_t global_num_layers,
+    int32_t pipeline_stage_flags,
+    int32_t compute_type,
+    double lora_scaling, double lr, double beta1, double beta2, double eps,
+    int64_t vocab_size, double rms_eps,
+    int64_t lora_rank,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str,
+    int32_t context_flags
+) {
+    try {
+        TORCH_CHECK(std::isfinite(lora_scaling) && lora_scaling > 0.0,
+            "fixed LoRA scaling must be finite and positive");
+        TORCH_CHECK(std::isfinite(lr) && lr >= 0.0,
+            "fixed Adam learning rate must be finite and non-negative");
+        TORCH_CHECK(std::isfinite(beta1) && beta1 >= 0.0 && beta1 < 1.0 &&
+                std::isfinite(beta2) && beta2 >= 0.0 && beta2 < 1.0,
+            "fixed Adam betas must be finite and in [0, 1)");
+        TORCH_CHECK(std::isfinite(eps) && eps >= 0.0,
+            "fixed Adam epsilon must be finite and non-negative");
+        auto* ctx = new TrainingContext();
+        ctx->context_sequence = ++g_context_sequence;
+        ctx->compute_type = static_cast<at::ScalarType>(compute_type);
+        ctx->lr = lr; ctx->beta1 = beta1; ctx->beta2 = beta2; ctx->eps = eps;
+        ctx->vocab_size = vocab_size; ctx->rms_eps = rms_eps;
+        ctx->fixed_optimizer_step = 0; ctx->lora_scaling = lora_scaling;
+        ctx->num_layers = num_layers;
+        ctx->global_layer_start = global_layer_start;
+        ctx->global_num_layers = global_num_layers;
+        ctx->is_first_pipeline_stage = (pipeline_stage_flags & 1) != 0;
+        ctx->is_last_pipeline_stage = (pipeline_stage_flags & 2) != 0;
+        const bool communicator_only =
+            num_layers == 0 && global_num_layers == 0;
+        TORCH_CHECK(communicator_only || num_layers > 0,
+            "pipeline stage must own at least one layer");
+        TORCH_CHECK(global_layer_start >= 0 && global_num_layers >= num_layers &&
+                global_layer_start <= global_num_layers - num_layers,
+            "invalid pipeline layer range [", global_layer_start, ", ",
+            global_layer_start + num_layers, ") for ", global_num_layers,
+            " global layers");
+        TORCH_CHECK(ctx->is_first_pipeline_stage == (global_layer_start == 0),
+            "first-stage flag does not match the global layer range");
+        TORCH_CHECK(ctx->is_last_pipeline_stage ==
+                (global_layer_start + num_layers == global_num_layers),
+            "last-stage flag does not match the global layer range");
+        ctx->base_tp_attention =
+            (context_flags & QWEN36_CONTEXT_BASE_TP_ATTENTION) != 0;
+        ctx->base_tp_mlp =
+            (context_flags & QWEN36_CONTEXT_BASE_TP_MLP) != 0;
+        ctx->vocab_parallel =
+            (context_flags & QWEN36_CONTEXT_VOCAB_PARALLEL) != 0;
+        ctx->sequence_parallel = env_enabled("QWEN36_SEQUENCE_PARALLEL");
+        ctx->has_mtp = false;
+        ctx->use_checkpoint = false; ctx->group_size = 4;
+        const char* tp_size_env = getenv("TP_SIZE");
+        if (!tp_size_env) tp_size_env = getenv("RUSTRAIN_TP_SIZE");
+        ctx->tp_world_size = tp_size_env ? atoi(tp_size_env) : 1;
+        TORCH_CHECK(ctx->tp_world_size > 0, "TP_SIZE must be positive");
+        const char* world_size_env = getenv("WORLD_SIZE");
+        const int configured_world_size = world_size_env ? atoi(world_size_env) : 1;
+        TORCH_CHECK(configured_world_size > 0, "WORLD_SIZE must be positive");
+        const bool data_parallel_requested =
+            (context_flags & QWEN36_CONTEXT_DATA_PARALLEL) != 0 ||
+            env_enabled("RUSTRAIN_DATA_PARALLEL");
+        const char* dp_size_env = getenv("DP_SIZE");
+        if (!dp_size_env) dp_size_env = getenv("RUSTRAIN_DP_SIZE");
+        int configured_dp_size = dp_size_env ? atoi(dp_size_env) : 1;
+        const char* ep_size_env = getenv("EP_SIZE");
+        if (!ep_size_env) ep_size_env = getenv("RUSTRAIN_EP_SIZE");
+        const int configured_ep_size = ep_size_env ? atoi(ep_size_env) : 1;
+        const char* pp_size_env = getenv("PP_SIZE");
+        if (!pp_size_env) pp_size_env = getenv("RUSTRAIN_PP_SIZE");
+        const char* cp_size_env = getenv("CP_SIZE");
+        if (!cp_size_env) cp_size_env = getenv("RUSTRAIN_CP_SIZE");
+        const int configured_pp_size = pp_size_env ? atoi(pp_size_env) : 1;
+        const int configured_cp_size = cp_size_env ? atoi(cp_size_env) : 1;
+        const bool expert_parallel_requested =
+            (context_flags & QWEN36_CONTEXT_EXPERT_PARALLEL) != 0 ||
+            configured_ep_size > 1;
+        const int64_t non_data_parallel_product =
+            static_cast<int64_t>(ctx->tp_world_size) * configured_cp_size *
+            configured_ep_size * configured_pp_size;
+        if (data_parallel_requested && !dp_size_env &&
+            non_data_parallel_product > 0 &&
+            configured_world_size % non_data_parallel_product == 0) {
+            configured_dp_size = static_cast<int>(
+                configured_world_size / non_data_parallel_product);
+        }
+        TORCH_CHECK(configured_pp_size > 0 && configured_cp_size > 0 &&
+                    configured_dp_size > 0 && configured_ep_size > 0,
+            "PP_SIZE, CP_SIZE, DP_SIZE, and EP_SIZE must be positive");
+        const int64_t configured_parallel_product =
+            static_cast<int64_t>(ctx->tp_world_size) * configured_cp_size *
+            configured_ep_size * configured_dp_size * configured_pp_size;
+        TORCH_CHECK(configured_parallel_product <= std::numeric_limits<int>::max() &&
+                configured_world_size == configured_parallel_product,
+            "native Qwen LoRA requires "
+            "WORLD_SIZE=TP_SIZE*CP_SIZE*EP_SIZE*DP_SIZE*PP_SIZE; ",
+            "TP_SIZE=", ctx->tp_world_size, " EP_SIZE=", configured_ep_size,
+            " DP_SIZE=", configured_dp_size, " PP_SIZE=", configured_pp_size,
+            " CP_SIZE=", configured_cp_size,
+            " WORLD_SIZE=", configured_world_size);
+        TORCH_CHECK(expert_parallel_requested == (configured_ep_size > 1),
+            "expert-parallel context flag must match EP_SIZE");
+        TORCH_CHECK(data_parallel_requested == (configured_dp_size > 1),
+            "data-parallel context flag must match DP_SIZE");
+        ctx->ep_world_size = configured_ep_size;
+        ctx->expert_parallel = expert_parallel_requested;
+        ctx->dp_world_size = configured_dp_size;
+        ctx->data_parallel = data_parallel_requested;
+        ctx->cp_world_size = configured_cp_size;
+        ctx->pp_world_size = configured_pp_size;
+        const char* rank_env = getenv("RANK");
+        const int global_rank = rank_env ? atoi(rank_env) : 0;
+        const char* tp_rank_env = getenv("RUSTRAIN_TP_RANK");
+        const char* cp_rank_env = getenv("RUSTRAIN_CP_RANK");
+        const char* ep_rank_env = getenv("RUSTRAIN_EP_RANK");
+        const char* dp_rank_env = getenv("RUSTRAIN_DP_RANK");
+        const char* pp_rank_env = getenv("RUSTRAIN_PP_RANK");
+        const char* parallel_order_env = getenv("RUSTRAIN_PARALLEL_ORDER");
+        if (!parallel_order_env) parallel_order_env = getenv("PARALLEL_ORDER");
+        if (parallel_order_env &&
+            std::strcmp(parallel_order_env, "tp-cp-ep-dp-pp") != 0) {
+            TORCH_CHECK(tp_rank_env && cp_rank_env && ep_rank_env &&
+                    dp_rank_env && pp_rank_env,
+                "custom parallel rank order requires explicit "
+                "RUSTRAIN_TP_RANK/RUSTRAIN_CP_RANK/RUSTRAIN_EP_RANK/"
+                "RUSTRAIN_DP_RANK/RUSTRAIN_PP_RANK coordinates");
+        }
+        ctx->tp_rank = tp_rank_env ? atoi(tp_rank_env)
+            : global_rank % ctx->tp_world_size;
+        ctx->cp_rank = cp_rank_env ? atoi(cp_rank_env)
+            : (global_rank / ctx->tp_world_size) % configured_cp_size;
+        ctx->ep_rank = ep_rank_env ? atoi(ep_rank_env)
+            : (global_rank / (ctx->tp_world_size * configured_cp_size)) %
+                configured_ep_size;
+        ctx->dp_rank = dp_rank_env ? atoi(dp_rank_env)
+            : (global_rank /
+                (ctx->tp_world_size * configured_cp_size * configured_ep_size)) %
+                configured_dp_size;
+        ctx->pp_rank = pp_rank_env ? atoi(pp_rank_env)
+            : global_rank /
+                (ctx->tp_world_size * configured_cp_size * configured_ep_size *
+                    configured_dp_size);
+        TORCH_CHECK(ctx->tp_rank >= 0 && ctx->tp_rank < ctx->tp_world_size &&
+                    ctx->cp_rank >= 0 && ctx->cp_rank < ctx->cp_world_size &&
+                    ctx->ep_rank >= 0 && ctx->ep_rank < ctx->ep_world_size &&
+                    ctx->dp_rank >= 0 && ctx->dp_rank < ctx->dp_world_size &&
+                    ctx->pp_rank >= 0 && ctx->pp_rank < ctx->pp_world_size,
+            "parallel rank coordinates are outside their configured axes");
+        if (!communicator_only) {
+            const int64_t expected_stage_start =
+                global_num_layers * ctx->pp_rank / ctx->pp_world_size;
+            const int64_t expected_stage_end =
+                global_num_layers * (ctx->pp_rank + 1) / ctx->pp_world_size;
+            TORCH_CHECK(global_layer_start == expected_stage_start &&
+                    global_layer_start + num_layers == expected_stage_end,
+                "pipeline stage layer ownership mismatch for PP rank ",
+                ctx->pp_rank, ": received [", global_layer_start, ", ",
+                global_layer_start + num_layers, "), expected [",
+                expected_stage_start, ", ", expected_stage_end, ")");
+        }
+        TORCH_CHECK(lora_rank > 0, "LoRA rank must be positive");
+        if (const char* mtp_scale = getenv("QWEN36_MTP_LOSS_SCALE")) {
+            ctx->mtp_loss_scale = std::strtod(mtp_scale, nullptr);
+        }
+
+        // Store weight pointers
+        auto** wp = reinterpret_cast<at::Tensor**>(weight_ptrs);
+        for (int64_t i = 0; i < num_weight_ptrs; i++) {
+            ctx->weight_ptrs.push_back(wp[i]);
+        }
+        auto* embed = reinterpret_cast<at::Tensor*>(embed_ptr);
+        auto* final_norm = reinterpret_cast<at::Tensor*>(final_norm_ptr);
+        auto* lm_head = reinterpret_cast<at::Tensor*>(lm_head_ptr);
+        TORCH_CHECK(ctx->is_first_pipeline_stage ? embed != nullptr : embed == nullptr,
+            "only the first pipeline stage may own the embedding weight");
+        TORCH_CHECK(ctx->is_last_pipeline_stage
+                ? final_norm != nullptr && lm_head != nullptr
+                : final_norm == nullptr && lm_head == nullptr,
+            "only the last pipeline stage may own final norm and LM-head weights");
+        ctx->embed_ptr.push_back(embed);
+        ctx->final_norm_ptr.push_back(final_norm);
+        ctx->lm_head_ptr.push_back(lm_head);
+        if (ctx->vocab_parallel) {
+            TORCH_CHECK(ctx->tp_world_size > 1,
+                "vocabulary parallelism requires TP_SIZE>1");
+            TORCH_CHECK(ctx->vocab_size > 0 &&
+                    ctx->vocab_size % ctx->tp_world_size == 0,
+                "vocab_size=", ctx->vocab_size,
+                " must be divisible by TP_SIZE=", ctx->tp_world_size);
+            ctx->local_vocab_size = ctx->vocab_size / ctx->tp_world_size;
+            if (embed) {
+                TORCH_CHECK(embed->dim() == 2 &&
+                        embed->size(0) == ctx->local_vocab_size,
+                    "vocabulary-parallel embedding has invalid local shape: ",
+                    embed->sizes(), " expected local vocab=", ctx->local_vocab_size);
+            }
+            if (lm_head) {
+                TORCH_CHECK(lm_head->dim() == 2 &&
+                        lm_head->size(0) == ctx->local_vocab_size,
+                    "vocabulary-parallel LM head has invalid local shape: ",
+                    lm_head->sizes(), " expected local vocab=", ctx->local_vocab_size);
+            }
+        } else {
+            ctx->local_vocab_size = ctx->vocab_size;
+        }
+
+        // Copy layer configs
+        auto* lcfgs = reinterpret_cast<LayerConfig*>(layer_configs_ptr);
+        for (int64_t i = 0; i < num_layers; i++) {
+            ctx->layer_configs.push_back(lcfgs[i]);
+        }
+        int64_t expected_weight_ptrs = 0;
+        for (const auto& cfg : ctx->layer_configs)
+            expected_weight_ptrs += weight_count_for_layer(cfg);
+        TORCH_CHECK(num_weight_ptrs == expected_weight_ptrs,
+            "pipeline stage received ", num_weight_ptrs,
+            " frozen layer weights, expected ", expected_weight_ptrs);
+        for (int64_t i = 0; i < num_weight_ptrs; ++i) {
+            TORCH_CHECK(ctx->weight_ptrs[i],
+                "pipeline stage received null frozen layer weight at local index ", i);
+        }
+
+        ctx->fused_gdn_ab_weights.resize(num_layers);
+        if (env_enabled("QWEN36_GDN_FUSED_AB_PROJECTION", true)) {
+            at::NoGradGuard guard;
+            int64_t weight_offset = 0;
+            for (int64_t layer = 0; layer < num_layers; ++layer) {
+                const auto& cfg = ctx->layer_configs[layer];
+                if (cfg.layer_type != 0) {
+                    auto* a = ctx->weight_ptrs[weight_offset + 4];
+                    auto* b = ctx->weight_ptrs[weight_offset + 5];
+                    TORCH_CHECK(a && b && a->dim() == 2 && b->dim() == 2 &&
+                            a->sizes() == b->sizes(),
+                        "GDN fused A/B projection requires matching matrix weights at layer ",
+                        layer);
+                    ctx->fused_gdn_ab_weights[layer] =
+                        at::cat({*a, *b}, 0).contiguous();
+                }
+                weight_offset += weight_count_for_layer(cfg);
+            }
+        }
+        ctx->fused_full_attention_qkv_weights.resize(num_layers);
+        if (env_enabled("QWEN36_FUSED_QKV", false)) {
+            at::NoGradGuard guard;
+            int64_t weight_offset = 0;
+            for (int64_t layer = 0; layer < num_layers; ++layer) {
+                const auto& cfg = ctx->layer_configs[layer];
+                if (cfg.layer_type == 0) {
+                    auto* q = ctx->weight_ptrs[weight_offset + 2];
+                    auto* k = ctx->weight_ptrs[weight_offset + 4];
+                    auto* v = ctx->weight_ptrs[weight_offset + 6];
+                    TORCH_CHECK(q && k && v && q->dim() == 2 &&
+                            k->dim() == 2 && v->dim() == 2 &&
+                            q->size(1) == k->size(1) &&
+                            q->size(1) == v->size(1),
+                        "fused full-attention QKV requires matching input dimensions at layer ",
+                        layer);
+                    ctx->fused_full_attention_qkv_weights[layer] =
+                        at::cat({*q, *k, *v}, 0).contiguous();
+                }
+                weight_offset += weight_count_for_layer(cfg);
+            }
+        }
+        ctx->fused_mlp_fc1_weights.resize(num_layers);
+        if (env_enabled("QWEN36_FUSED_MLP_FC1", false)) {
+            at::NoGradGuard guard;
+            int64_t weight_offset = 0;
+            for (int64_t layer = 0; layer < num_layers; ++layer) {
+                const auto& cfg = ctx->layer_configs[layer];
+                const int64_t mlp_start = cfg.layer_type == 0 ? 8 : 11;
+                const int64_t gate_offset = cfg.num_experts > 0
+                    ? mlp_start + 2 : mlp_start;
+                const int64_t up_offset = gate_offset + 1;
+                auto* gate = ctx->weight_ptrs[weight_offset + gate_offset];
+                auto* up = ctx->weight_ptrs[weight_offset + up_offset];
+                TORCH_CHECK(gate && up && gate->dim() == 2 && up->dim() == 2 &&
+                        gate->sizes() == up->sizes(),
+                    "fused MLP FC1 requires matching gate/up matrix weights at layer ",
+                    layer);
+                ctx->fused_mlp_fc1_weights[layer] =
+                    at::cat({*gate, *up}, 0).contiguous();
+                weight_offset += weight_count_for_layer(cfg);
+            }
+        }
+
+        if (ctx->base_tp_attention) {
+            TORCH_CHECK(ctx->tp_world_size > 1,
+                "base attention TP requires TP_SIZE>1");
+            int64_t weight_offset = 0;
+            for (int64_t layer = 0; layer < num_layers; ++layer) {
+                const auto& cfg = ctx->layer_configs[layer];
+                if (cfg.layer_type == 0) {
+                    TORCH_CHECK(cfg.num_heads > 0 && cfg.num_kv_heads > 0 &&
+                            cfg.head_dim > 0,
+                        "invalid full-attention head configuration at layer ", layer);
+                    TORCH_CHECK(cfg.num_heads % cfg.num_kv_heads == 0 &&
+                            cfg.num_heads % ctx->tp_world_size == 0 &&
+                            cfg.num_kv_heads % ctx->tp_world_size == 0,
+                        "full-attention heads must preserve GQA groups and be divisible by TP_SIZE at layer ",
+                        layer);
+                    const int64_t rotary_dim = static_cast<int64_t>(
+                        cfg.head_dim * cfg.partial_rotary_factor);
+                    TORCH_CHECK(rotary_dim >= 0 && rotary_dim <= cfg.head_dim &&
+                            rotary_dim % 2 == 0,
+                        "full-attention rotary dimension must be even and within head_dim at layer ",
+                        layer, ": rotary_dim=", rotary_dim,
+                        " head_dim=", cfg.head_dim);
+                    auto* q = ctx->weight_ptrs[weight_offset + 2];
+                    auto* k = ctx->weight_ptrs[weight_offset + 4];
+                    auto* v = ctx->weight_ptrs[weight_offset + 6];
+                    auto* o = ctx->weight_ptrs[weight_offset + 7];
+                    TORCH_CHECK(q && k && v && o && q->dim() == 2 &&
+                            k->dim() == 2 && v->dim() == 2 && o->dim() == 2,
+                        "base full-attention TP requires matrix Q/K/V/O weights at layer ",
+                        layer);
+                    const int64_t local_heads = cfg.num_heads / ctx->tp_world_size;
+                    const int64_t local_kv_heads = cfg.num_kv_heads / ctx->tp_world_size;
+                    TORCH_CHECK(q->size(0) == local_heads * cfg.head_dim * 2 &&
+                            k->size(0) == local_kv_heads * cfg.head_dim &&
+                            v->size(0) == local_kv_heads * cfg.head_dim &&
+                            o->size(1) == local_heads * cfg.head_dim &&
+                            q->size(1) == k->size(1) && q->size(1) == v->size(1) &&
+                            q->size(1) == o->size(0),
+                        "base full-attention TP received inconsistent local weight shapes at layer ",
+                        layer, ": q=", q->sizes(), " k=", k->sizes(),
+                        " v=", v->sizes(), " o=", o->sizes());
+                } else {
+                    TORCH_CHECK(cfg.num_k_heads > 0 && cfg.num_v_heads > 0 &&
+                            cfg.key_dim == 128 && cfg.val_dim == 128 &&
+                            cfg.conv_kernel > 0,
+                        "invalid linear-attention configuration at layer ", layer);
+                    TORCH_CHECK(cfg.num_v_heads % cfg.num_k_heads == 0 &&
+                            cfg.num_k_heads % ctx->tp_world_size == 0 &&
+                            cfg.num_v_heads % ctx->tp_world_size == 0,
+                        "linear-attention heads must preserve value-head groups and be divisible by TP_SIZE at layer ",
+                        layer);
+                    auto* qkv = ctx->weight_ptrs[weight_offset + 2];
+                    auto* z = ctx->weight_ptrs[weight_offset + 3];
+                    auto* a = ctx->weight_ptrs[weight_offset + 4];
+                    auto* b = ctx->weight_ptrs[weight_offset + 5];
+                    auto* a_log = ctx->weight_ptrs[weight_offset + 6];
+                    auto* dt_bias = ctx->weight_ptrs[weight_offset + 7];
+                    auto* conv = ctx->weight_ptrs[weight_offset + 8];
+                    auto* norm = ctx->weight_ptrs[weight_offset + 9];
+                    auto* out = ctx->weight_ptrs[weight_offset + 10];
+                    TORCH_CHECK(qkv && z && a && b && a_log && dt_bias &&
+                            conv && norm && out,
+                        "base linear-attention TP received null weights at layer ", layer);
+                    const int64_t local_k_heads =
+                        cfg.num_k_heads / ctx->tp_world_size;
+                    const int64_t local_v_heads =
+                        cfg.num_v_heads / ctx->tp_world_size;
+                    const int64_t local_q = local_k_heads * cfg.key_dim;
+                    const int64_t local_v = local_v_heads * cfg.val_dim;
+                    const int64_t local_qkv = local_q * 2 + local_v;
+                    TORCH_CHECK(qkv->dim() == 2 &&
+                            qkv->size(0) == local_qkv,
+                        "base linear-attention TP QKV shape mismatch at layer ",
+                        layer, ": ", qkv->sizes(), " expected rows=", local_qkv);
+                    const int64_t hidden_size = qkv->size(1);
+                    TORCH_CHECK(z->dim() == 2 && z->size(0) == local_v &&
+                            z->size(1) == hidden_size &&
+                            a->dim() == 2 && a->size(0) == local_v_heads &&
+                            a->size(1) == hidden_size &&
+                            b->dim() == 2 && b->sizes() == a->sizes(),
+                        "base linear-attention TP Z/A/B shapes mismatch at layer ", layer,
+                        ": z=", z->sizes(), " a=", a->sizes(), " b=", b->sizes());
+                    TORCH_CHECK(a_log->dim() == 1 &&
+                            a_log->size(0) == local_v_heads &&
+                            dt_bias->dim() == 1 &&
+                            dt_bias->size(0) == local_v_heads,
+                        "base linear-attention TP A_log/dt_bias shapes mismatch at layer ",
+                        layer, ": A_log=", a_log->sizes(),
+                        " dt_bias=", dt_bias->sizes());
+                    TORCH_CHECK(conv->dim() == 3 &&
+                            conv->size(0) == local_qkv && conv->size(1) == 1 &&
+                            conv->size(2) == cfg.conv_kernel,
+                        "base linear-attention TP depthwise-conv shape mismatch at layer ",
+                        layer, ": ", conv->sizes());
+                    TORCH_CHECK(norm->dim() == 1 && norm->size(0) == cfg.val_dim,
+                        "base linear-attention TP norm must remain replicated at layer ",
+                        layer, ": ", norm->sizes());
+                    TORCH_CHECK(out->dim() == 2 && out->size(0) == hidden_size &&
+                            out->size(1) == local_v,
+                        "base linear-attention TP output shape mismatch at layer ",
+                        layer, ": ", out->sizes(), " expected [", hidden_size,
+                        ", ", local_v, "]");
+                }
+                weight_offset += weight_count_for_layer(cfg);
+            }
+        }
+
+        // Build target layer set
+        std::set<int64_t> target_set;
+        const bool all_target_layers = !target_layers || num_target_layers == 0;
+        ctx->fixed_all_target_layers = all_target_layers;
+        if (target_layers && num_target_layers > 0) {
+            for (int64_t j = 0; j < num_target_layers; j++) {
+                TORCH_CHECK(target_layers[j] >= 0 && target_layers[j] < global_num_layers,
+                    "LoRA target layer out of range: ", target_layers[j],
+                    " for model with ", global_num_layers, " layers");
+                ctx->fixed_target_layers_global.insert(target_layers[j]);
+                if (target_layers[j] >= global_layer_start &&
+                        target_layers[j] < global_layer_start + num_layers) {
+                    target_set.insert(target_layers[j] - global_layer_start);
+                }
+            }
+        }
+
+        std::set<std::string> target_modules;
+        if (target_modules_str && target_modules_str[0] != '\0') {
+            std::stringstream ss(target_modules_str);
+            std::string item;
+            while (std::getline(ss, item, ',')) {
+                if (!item.empty()) target_modules.insert(item);
+            }
+        }
+        ctx->fixed_target_modules = target_modules;
+        for (const auto& name : target_modules) {
+            TORCH_CHECK(
+                name == "q_proj" || name == "k_proj" || name == "v_proj" ||
+                name == "o_proj" || name == "in_proj_qkv" ||
+                name == "in_proj_z" || name == "in_proj_a" ||
+                name == "in_proj_b" || name == "out_proj" ||
+                name == "gate_proj" || name == "up_proj" ||
+                name == "down_proj" || name == "shared_gate_proj" ||
+                name == "shared_up_proj" || name == "shared_down_proj" ||
+                name == "experts_gate_up_proj" || name == "experts_down_proj",
+                "unsupported native Qwen LoRA target module: ", name,
+                "; supported routed expert targets are experts_gate_up_proj/experts_down_proj");
+        }
+        for (const auto& name : target_modules) {
+            bool resolved = false;
+            for (const auto& layer_cfg : ctx->layer_configs) {
+                auto table = lora_projection_table(layer_cfg);
+                for (int64_t k = 0; k < table.count; ++k) {
+                    if (name == table.entries[k].name) {
+                        resolved = true;
+                        break;
+                    }
+                }
+                if (resolved) break;
+            }
+            TORCH_CHECK(resolved || global_num_layers != num_layers,
+                "LoRA target module does not exist in this model: ", name);
+        }
+        const int64_t local_lora_rank = local_lora_rank_for_active_targets(
+            ctx, lora_rank, target_set, target_modules,
+            all_target_layers,
+            /*empty_modules_mean_attention_only=*/false, "fixed");
+
+        // Create fixed positional slots for every layer so layer offsets stay
+        // stable. Inactive slots are zero tensors without grad and are skipped
+        // by cache construction/Adam; this preserves the existing FFI export
+        // indexing while honoring target_modules exactly.
+        int64_t offset = 0;
+        for (int64_t i = 0; i < num_layers; i++) {
+            auto projection_table = lora_projection_table(ctx->layer_configs[i]);
+            int64_t lora_count = projection_table.count;
+            ctx->lora_layer_offset.push_back(offset);
+
+            // Get base weight shapes from the weight pointers
+            int64_t w_offset = 0;
+            for (int64_t j = 0; j < i; j++)
+                w_offset += weight_count_for_layer(ctx->layer_configs[j]);
+
+            for (int64_t k = 0; k < projection_table.count; ++k) {
+                const auto& projection = projection_table.entries[k];
+                auto* base = ctx->weight_ptrs[w_offset + projection.weight_index];
+                TORCH_CHECK(base, "null LoRA base projection: layer=", i,
+                    " module=", projection.name);
+                TORCH_CHECK(
+                    (!projection.grouped_expert && base->dim() == 2)
+                        || (projection.grouped_expert && base->dim() == 3),
+                    "LoRA projection rank mismatch: layer=", i,
+                    " module=", projection.name, " base_dim=", base->dim());
+                bool active = (all_target_layers || target_set.find(i) != target_set.end()) &&
+                    (target_modules.empty() || target_modules.count(projection.name) > 0);
+                auto opts = at::TensorOptions().dtype(ctx->compute_type).device(base->device());
+                at::Tensor a, b;
+                if (!active) {
+                    // Stable slot index without allocating potentially hundreds
+                    // of MB for inactive expert-local tensors.
+                    a = at::zeros({}, opts);
+                    b = at::zeros({}, opts);
+                } else if (projection.grouped_expert) {
+                    int64_t experts = base->size(0);
+                    int64_t out_f = base->size(1), in_f = base->size(2);
+                    const auto layout = lora_tp_layout(ctx, i, k);
+                    if (layout == LoraTpLayout::ColumnParallel ||
+                        layout == LoraTpLayout::RowParallel) {
+                        a = at::randn({experts, lora_rank, in_f}, opts) * 0.01;
+                        b = at::zeros({experts, out_f, lora_rank}, opts);
+                    } else {
+                        a = initialize_lora_a(ctx, opts, experts, lora_rank, in_f);
+                        b = at::zeros({experts, out_f, local_lora_rank}, opts);
+                    }
+                } else {
+                    int64_t out_f = base->size(0), in_f = base->size(1);
+                    const auto layout = lora_tp_layout(ctx, i, k);
+                    if (layout == LoraTpLayout::ColumnParallel) {
+                        a = at::randn({lora_rank, in_f}, opts) * 0.01;
+                        b = at::zeros({out_f, lora_rank}, opts);
+                    } else if (layout == LoraTpLayout::RowParallel) {
+                        a = at::randn({lora_rank, in_f}, opts) * 0.01;
+                        b = at::zeros({out_f, lora_rank}, opts);
+                    } else {
+                        a = initialize_lora_a(ctx, opts, 0, lora_rank, in_f);
+                        b = at::zeros({out_f, local_lora_rank}, opts);
+                    }
+                }
+                a.set_requires_grad(active);
+                b.set_requires_grad(active);
+                ctx->grad_accum_a.push_back(at::Tensor());
+                ctx->grad_accum_b.push_back(at::Tensor());
+                ctx->lora_a.push_back(std::move(a));
+                ctx->lora_b.push_back(std::move(b));
+                ctx->lora_active.push_back(active ? 1 : 0);
+                auto prefix = "layers." +
+                    std::to_string(global_layer_start + i) + "." + projection.name;
+                ctx->lora_names.push_back(prefix + ".lora_A.weight");
+                ctx->lora_names.push_back(prefix + ".lora_B.weight");
+            }
+            offset += lora_count;
+        }
+        bind_fixed_lora_gradient_slab(ctx);
+
+        // Initialize Adam state (FP32 for numerical stability, even if params are BF16)
+        for (size_t i = 0; i < ctx->lora_a.size(); i++) {
+            auto opts_f32 = at::TensorOptions().dtype(at::kFloat).device(ctx->lora_a[i].device());
+            ctx->adam_m.push_back(at::zeros(ctx->lora_a[i].sizes(), opts_f32));
+            ctx->adam_m.push_back(at::zeros(ctx->lora_b[i].sizes(), opts_f32));
+            ctx->adam_v.push_back(at::zeros(ctx->lora_a[i].sizes(), opts_f32));
+            ctx->adam_v.push_back(at::zeros(ctx->lora_b[i].sizes(), opts_f32));
+        }
+
+        fprintf(stderr,
+            "[q36_ctx] created: layers=[%ld,%ld)/%ld, %ld LoRA params, %ld Adam states\n",
+            (long)global_layer_start, (long)(global_layer_start + num_layers),
+            (long)global_num_layers, (long)ctx->lora_a.size(),
+            (long)ctx->adam_m.size());
+        return ctx;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] create FAILED: %s\n", e.what());
+        return nullptr;
+    }
+}
+
+__attribute__((visibility("default"))) void* qwen36_create_training_context(
+    void** weight_ptrs, int64_t num_weight_ptrs,
+    void* embed_ptr, void* final_norm_ptr, void* lm_head_ptr,
+    void* layer_configs_ptr, int64_t num_layers,
+    int32_t compute_type,
+    double lora_scaling, double lr, double beta1, double beta2, double eps,
+    int64_t vocab_size, double rms_eps,
+    int64_t lora_rank,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str
+) {
+    return qwen36_create_training_context_impl(
+        weight_ptrs, num_weight_ptrs, embed_ptr, final_norm_ptr, lm_head_ptr,
+        layer_configs_ptr, num_layers, 0, num_layers, 3, compute_type,
+        lora_scaling, lr, beta1,
+        beta2, eps, vocab_size, rms_eps, lora_rank, target_layers,
+        num_target_layers, target_modules_str, 0);
+}
+
+__attribute__((visibility("default"))) void* qwen36_create_training_context_ex(
+    void** weight_ptrs, int64_t num_weight_ptrs,
+    void* embed_ptr, void* final_norm_ptr, void* lm_head_ptr,
+    void* layer_configs_ptr, int64_t num_layers,
+    int32_t compute_type,
+    double lora_scaling, double lr, double beta1, double beta2, double eps,
+    int64_t vocab_size, double rms_eps,
+    int64_t lora_rank,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str, int32_t context_flags
+) {
+    return qwen36_create_training_context_impl(
+        weight_ptrs, num_weight_ptrs, embed_ptr, final_norm_ptr, lm_head_ptr,
+        layer_configs_ptr, num_layers, 0, num_layers, 3, compute_type,
+        lora_scaling, lr, beta1,
+        beta2, eps, vocab_size, rms_eps, lora_rank, target_layers,
+        num_target_layers, target_modules_str, context_flags);
+}
+
+__attribute__((visibility("default"))) void* qwen36_create_training_context_v2(
+    void** weight_ptrs, int64_t num_weight_ptrs,
+    void* embed_ptr, void* final_norm_ptr, void* lm_head_ptr,
+    void* layer_configs_ptr, int64_t num_layers,
+    int64_t global_layer_start, int64_t global_num_layers,
+    int32_t pipeline_stage_flags,
+    int32_t compute_type,
+    double lora_scaling, double lr, double beta1, double beta2, double eps,
+    int64_t vocab_size, double rms_eps,
+    int64_t lora_rank,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str, int32_t context_flags
+) {
+    return qwen36_create_training_context_impl(
+        weight_ptrs, num_weight_ptrs, embed_ptr, final_norm_ptr, lm_head_ptr,
+        layer_configs_ptr, num_layers, global_layer_start, global_num_layers,
+        pipeline_stage_flags, compute_type, lora_scaling, lr, beta1, beta2,
+        eps, vocab_size, rms_eps, lora_rank, target_layers,
+        num_target_layers, target_modules_str, context_flags);
+}
+
+__attribute__((visibility("default"))) int32_t qwen36_set_base_tp_mlp(
+    void* ctx_ptr, int32_t enabled
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "null training context");
+        if (!enabled) {
+            TORCH_CHECK(!ctx->base_tp_mlp,
+                "base MLP TP cannot be disabled after enablement because "
+                "the context owns TP-sharded weights");
+            return 0;
+        }
+        const bool preconfigured = ctx->base_tp_mlp;
+        TORCH_CHECK(ctx->tp_world_size > 1,
+            "base MLP TP requires TP_SIZE>1");
+        TORCH_CHECK(!ctx->has_mtp || preconfigured,
+            "base MLP TP must be selected when the context is created before "
+            "MTP weights are configured");
+        if (!preconfigured) {
+            for (const auto& adapter : ctx->adapters) {
+                TORCH_CHECK(!adapter.target_modules.empty(),
+                    "base MLP TP must be selected when the context is created "
+                    "before adding a dynamic adapter that targets all modules");
+                for (const auto& name : adapter.target_modules) {
+                    TORCH_CHECK(!is_mlp_lora_target(name),
+                        "base MLP TP must be selected when the context is created "
+                        "before adding dynamic MLP LoRA target ", name);
+                }
+            }
+        }
+
+        int64_t weight_offset = 0;
+        for (int64_t layer = 0; layer < ctx->num_layers; ++layer) {
+            const auto& cfg = ctx->layer_configs[layer];
+            const int64_t mlp_start = cfg.layer_type == 0 ? 8 : 11;
+            if (cfg.num_experts == 0) {
+                TORCH_CHECK(cfg.intermediate_size > 0 &&
+                        cfg.intermediate_size % ctx->tp_world_size == 0,
+                    "dense intermediate_size must be divisible by TP_SIZE");
+                const int64_t local_intermediate =
+                    cfg.intermediate_size / ctx->tp_world_size;
+                auto* gate = ctx->weight_ptrs[weight_offset + mlp_start];
+                auto* up = ctx->weight_ptrs[weight_offset + mlp_start + 1];
+                auto* down = ctx->weight_ptrs[weight_offset + mlp_start + 2];
+                TORCH_CHECK(gate && up && down && gate->dim() == 2 &&
+                        up->dim() == 2 && down->dim() == 2,
+                    "base dense MLP TP requires matrix gate/up/down weights");
+                TORCH_CHECK(gate->size(0) == local_intermediate &&
+                        up->size(0) == local_intermediate &&
+                        down->size(1) == local_intermediate &&
+                        gate->size(1) == up->size(1) &&
+                        gate->size(1) == down->size(0),
+                    "base dense MLP TP received inconsistent local weight shapes: gate=",
+                    gate->sizes(), " up=", up->sizes(), " down=", down->sizes(),
+                    " expected local intermediate=", local_intermediate);
+            } else {
+                TORCH_CHECK(cfg.moe_intermediate > 0 &&
+                        cfg.moe_intermediate % ctx->tp_world_size == 0,
+                    "routed expert intermediate must be divisible by TP_SIZE");
+                const int64_t local_intermediate =
+                    cfg.moe_intermediate / ctx->tp_world_size;
+                auto* shared_gate = ctx->weight_ptrs[weight_offset + mlp_start + 2];
+                auto* shared_up = ctx->weight_ptrs[weight_offset + mlp_start + 3];
+                auto* shared_down = ctx->weight_ptrs[weight_offset + mlp_start + 4];
+                auto* experts_gate_up =
+                    ctx->weight_ptrs[weight_offset + mlp_start + 5];
+                auto* experts_down =
+                    ctx->weight_ptrs[weight_offset + mlp_start + 6];
+                TORCH_CHECK(shared_gate && shared_up && shared_down &&
+                        experts_gate_up && experts_down &&
+                        shared_gate->dim() == 2 && shared_up->dim() == 2 &&
+                        shared_down->dim() == 2 &&
+                        experts_gate_up->dim() == 3 && experts_down->dim() == 3,
+                    "base expert TP requires matrix shared weights and rank-3 "
+                    "routed expert weights");
+                TORCH_CHECK(shared_gate->size(0) > 0 &&
+                        shared_gate->sizes() == shared_up->sizes() &&
+                        shared_down->size(1) == shared_gate->size(0) &&
+                        shared_down->size(0) == shared_gate->size(1),
+                    "base shared expert TP received inconsistent local weight shapes: gate=",
+                    shared_gate->sizes(), " up=", shared_up->sizes(),
+                    " down=", shared_down->sizes());
+                TORCH_CHECK(experts_gate_up->size(0) == cfg.expert_count &&
+                        experts_down->size(0) == cfg.expert_count &&
+                        experts_gate_up->size(1) == 2 * local_intermediate &&
+                        experts_down->size(2) == local_intermediate &&
+                        experts_gate_up->size(2) == shared_gate->size(1) &&
+                        experts_down->size(1) == shared_gate->size(1),
+                    "base routed expert TP received inconsistent local weight shapes: gate_up=",
+                    experts_gate_up->sizes(), " down=", experts_down->sizes(),
+                    " expected local intermediate=", local_intermediate,
+                    " local experts=", cfg.expert_count);
+            }
+
+            if (!preconfigured) {
+                const int64_t lora_offset = ctx->lora_layer_offset[layer];
+                auto projections = lora_projection_table(cfg);
+                for (int64_t pair = 0; pair < projections.count; ++pair) {
+                    if (projections.entries[pair].segment == LoraSegment::Mlp) {
+                        TORCH_CHECK(!ctx->lora_active[lora_offset + pair],
+                            "base MLP TP must be selected when the context is "
+                            "created before enabling fixed MLP LoRA target ",
+                            projections.entries[pair].name);
+                    }
+                }
+            }
+            weight_offset += weight_count_for_layer(cfg);
+        }
+        ctx->base_tp_mlp = true;
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] set_base_tp_mlp FAILED: %s\n", e.what());
+        return -1;
+    }
+}
+
+__attribute__((visibility("default"))) int32_t qwen36_set_max_grad_norm(
+    void* ctx_ptr, double max_grad_norm
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "null training context");
+        TORCH_CHECK(std::isfinite(max_grad_norm) && max_grad_norm >= 0.0 &&
+                max_grad_norm <= static_cast<double>(std::numeric_limits<float>::max()),
+            "max gradient norm must be finite, non-negative, and representable as FP32");
+        ctx->max_grad_norm = max_grad_norm;
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] set_max_grad_norm FAILED: %s\n", e.what());
+        return -1;
+    }
+}
+
+__attribute__((visibility("default"))) int32_t qwen36_set_router_aux_loss_coef(
+    void* ctx_ptr, double coefficient
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "null training context");
+        TORCH_CHECK(std::isfinite(coefficient) && coefficient >= 0.0,
+            "router auxiliary loss coefficient must be finite and non-negative");
+        TORCH_CHECK(!ctx->accumulation_active && !ctx->pp_window.active,
+            "router auxiliary loss coefficient cannot change during training");
+        TORCH_CHECK(ctx->adapters.empty(),
+            "router auxiliary loss cannot be enabled after dynamic adapters are registered");
+        TORCH_CHECK(!ctx->has_mtp || coefficient == 0.0,
+            "router auxiliary loss with MTP is not yet supported");
+        ctx->router_aux_loss_coef = coefficient;
+        return 0;
+    } catch (const std::exception& error) {
+        fprintf(stderr, "[q36] set router aux loss FAILED: %s\n", error.what());
+        return -1;
+    }
+}
+
+static void validate_mtp_prediction_layer_layout(
+    const TrainingContext* ctx,
+    const LayerConfig& cfg,
+    at::Tensor** weights,
+    int64_t weight_offset,
+    int64_t layer_index,
+    int64_t hidden_size
+) {
+    TORCH_CHECK(cfg.layer_type == 0,
+        "MTP prediction layer ", layer_index,
+        " must use full attention");
+    auto* input_norm = weights[weight_offset];
+    auto* post_norm = weights[weight_offset + 1];
+    TORCH_CHECK(input_norm->dim() == 1 && input_norm->size(0) == hidden_size &&
+            post_norm->dim() == 1 && post_norm->size(0) == hidden_size,
+        "MTP prediction layer ", layer_index,
+        " normalization weights must have shape [hidden]");
+
+    TORCH_CHECK(cfg.num_heads > 0 && cfg.num_kv_heads > 0 &&
+            cfg.head_dim > 0 && cfg.num_heads % cfg.num_kv_heads == 0,
+        "MTP prediction layer ", layer_index,
+        " has an invalid full-attention head configuration");
+    const int64_t attention_partitions =
+        ctx->base_tp_attention ? ctx->tp_world_size : 1;
+    TORCH_CHECK(cfg.num_heads % attention_partitions == 0 &&
+            cfg.num_kv_heads % attention_partitions == 0,
+        "MTP prediction layer ", layer_index,
+        " attention heads must be divisible by its TP partition count");
+    const int64_t rotary_dim = static_cast<int64_t>(
+        cfg.head_dim * cfg.partial_rotary_factor);
+    TORCH_CHECK(rotary_dim >= 0 && rotary_dim <= cfg.head_dim &&
+            rotary_dim % 2 == 0,
+        "MTP prediction layer ", layer_index,
+        " has an invalid rotary dimension");
+    const int64_t local_heads = cfg.num_heads / attention_partitions;
+    const int64_t local_kv_heads = cfg.num_kv_heads / attention_partitions;
+    auto* q = weights[weight_offset + 2];
+    auto* q_norm = weights[weight_offset + 3];
+    auto* k = weights[weight_offset + 4];
+    auto* k_norm = weights[weight_offset + 5];
+    auto* v = weights[weight_offset + 6];
+    auto* o = weights[weight_offset + 7];
+    TORCH_CHECK(q->dim() == 2 && k->dim() == 2 && v->dim() == 2 &&
+            o->dim() == 2 && q_norm->dim() == 1 && k_norm->dim() == 1,
+        "MTP prediction layer ", layer_index,
+        " requires matrix Q/K/V/O and vector Q/K norm weights");
+    TORCH_CHECK(q->size(0) == local_heads * cfg.head_dim * 2 &&
+            k->size(0) == local_kv_heads * cfg.head_dim &&
+            v->size(0) == local_kv_heads * cfg.head_dim &&
+            q->size(1) == hidden_size && k->size(1) == hidden_size &&
+            v->size(1) == hidden_size && o->size(0) == hidden_size &&
+            o->size(1) == local_heads * cfg.head_dim &&
+            q_norm->size(0) == cfg.head_dim &&
+            k_norm->size(0) == cfg.head_dim,
+        "MTP prediction layer ", layer_index,
+        " received inconsistent local attention weights: q=", q->sizes(),
+        " k=", k->sizes(), " v=", v->sizes(), " o=", o->sizes(),
+        " attention partitions=", attention_partitions);
+
+    const int64_t mlp_partitions = ctx->base_tp_mlp
+        ? ctx->tp_world_size : 1;
+    if (cfg.num_experts <= 0) {
+        TORCH_CHECK(cfg.intermediate_size > 0 &&
+                cfg.intermediate_size % mlp_partitions == 0,
+            "MTP prediction layer ", layer_index,
+            " dense intermediate size must be divisible by its TP partition count");
+        const int64_t local_intermediate =
+            cfg.intermediate_size / mlp_partitions;
+        auto* gate = weights[weight_offset + 8];
+        auto* up = weights[weight_offset + 9];
+        auto* down = weights[weight_offset + 10];
+        TORCH_CHECK(gate->dim() == 2 && up->dim() == 2 && down->dim() == 2 &&
+                gate->size(0) == local_intermediate &&
+                up->size(0) == local_intermediate &&
+                gate->size(1) == hidden_size && up->size(1) == hidden_size &&
+                down->size(0) == hidden_size &&
+                down->size(1) == local_intermediate,
+            "MTP prediction layer ", layer_index,
+            " received inconsistent local dense MLP weights: gate=",
+            gate->sizes(), " up=", up->sizes(), " down=", down->sizes(),
+            " MLP partitions=", mlp_partitions);
+        return;
+    }
+
+    TORCH_CHECK(cfg.moe_intermediate > 0 &&
+            cfg.moe_intermediate % mlp_partitions == 0,
+        "MTP prediction layer ", layer_index,
+        " routed intermediate size must be divisible by its TP partition count");
+    TORCH_CHECK(cfg.top_k > 0 && cfg.top_k <= cfg.num_experts,
+        "MTP prediction layer ", layer_index,
+        " has an invalid routed-expert top-k");
+    TORCH_CHECK(ctx->ep_world_size > 0 &&
+            cfg.num_experts % ctx->ep_world_size == 0 &&
+            cfg.expert_count == cfg.num_experts / ctx->ep_world_size &&
+            cfg.expert_start == ctx->ep_rank * cfg.expert_count,
+        "MTP prediction layer ", layer_index,
+        " requires equal rank-contiguous expert ownership: global_experts=",
+        cfg.num_experts, " EP_SIZE=", ctx->ep_world_size,
+        " EP_RANK=", ctx->ep_rank, " expert_start=", cfg.expert_start,
+        " expert_count=", cfg.expert_count);
+    const int64_t local_intermediate =
+        cfg.moe_intermediate / mlp_partitions;
+    auto* router = weights[weight_offset + 8];
+    auto* shared_router = weights[weight_offset + 9];
+    auto* shared_gate = weights[weight_offset + 10];
+    auto* shared_up = weights[weight_offset + 11];
+    auto* shared_down = weights[weight_offset + 12];
+    auto* experts_gate_up = weights[weight_offset + 13];
+    auto* experts_down = weights[weight_offset + 14];
+    TORCH_CHECK(router->dim() == 2 &&
+            router->sizes() == at::IntArrayRef({cfg.num_experts, hidden_size}) &&
+            shared_router->dim() == 2 &&
+            shared_router->sizes() == at::IntArrayRef({1, hidden_size}),
+        "MTP prediction layer ", layer_index,
+        " router weights must remain replicated matrices over hidden");
+    TORCH_CHECK(shared_gate->dim() == 2 && shared_up->dim() == 2 &&
+            shared_down->dim() == 2 &&
+            shared_gate->sizes() == shared_up->sizes() &&
+            shared_gate->size(0) > 0 && shared_gate->size(1) == hidden_size &&
+            shared_down->size(0) == hidden_size &&
+            shared_down->size(1) == shared_gate->size(0),
+        "MTP prediction layer ", layer_index,
+        " received inconsistent local shared-expert weights");
+    TORCH_CHECK(experts_gate_up->dim() == 3 && experts_down->dim() == 3 &&
+            experts_gate_up->size(0) == cfg.expert_count &&
+            experts_down->size(0) == cfg.expert_count &&
+            experts_gate_up->size(1) == 2 * local_intermediate &&
+            experts_gate_up->size(2) == hidden_size &&
+            experts_down->size(1) == hidden_size &&
+            experts_down->size(2) == local_intermediate,
+        "MTP prediction layer ", layer_index,
+        " received inconsistent local routed-expert weights: gate_up=",
+        experts_gate_up->sizes(), " down=", experts_down->sizes(),
+        " MLP partitions=", mlp_partitions);
+}
+
+// Set MTP weights on an existing training context.
+// Called after create_training_context if MTP is enabled.
+__attribute__((visibility("default"))) int32_t qwen36_set_mtp_weights(
+    void* ctx_ptr,
+    void* mtp_fc_ptr,
+    void* mtp_pre_fc_norm_emb_ptr,
+    void* mtp_pre_fc_norm_hidden_ptr,
+    void* mtp_norm_ptr,
+    void** mtp_layer_weight_ptrs, int64_t num_mtp_layer_weights,
+    void* mtp_layer_configs_ptr, int64_t num_mtp_layers
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "null training context");
+        TORCH_CHECK(ctx->router_aux_loss_coef == 0.0,
+            "MTP cannot be enabled with router auxiliary loss");
+        TORCH_CHECK(!ctx->sequence_parallel,
+            "MTP currently requires replicated sequence state");
+        TORCH_CHECK(ctx->tp_world_size == 1 ||
+                (ctx->tp_world_size == 2 &&
+                    ctx->base_tp_attention && ctx->base_tp_mlp),
+            "TP MTP requires TP_SIZE=2 with attention and MLP TP selected "
+            "when the context is created");
+        TORCH_CHECK(!ctx->has_mtp, "MTP weights are already configured");
+        TORCH_CHECK(!ctx->accumulation_active && !ctx->pp_window.active,
+            "MTP weights cannot be configured during an active training window");
+        TORCH_CHECK(std::isfinite(ctx->mtp_loss_scale) &&
+                ctx->mtp_loss_scale >= 0.0,
+            "MTP loss scale must be finite and non-negative");
+        TORCH_CHECK(mtp_fc_ptr && mtp_pre_fc_norm_emb_ptr &&
+                mtp_pre_fc_norm_hidden_ptr && mtp_norm_ptr,
+            "MTP requires projection and normalization tensors");
+        TORCH_CHECK(num_mtp_layers > 0 && mtp_layer_configs_ptr,
+            "MTP requires at least one configured prediction layer");
+
+        auto* mtp_fc = reinterpret_cast<at::Tensor*>(mtp_fc_ptr);
+        auto* pre_norm_emb = reinterpret_cast<at::Tensor*>(
+            mtp_pre_fc_norm_emb_ptr);
+        auto* pre_norm_hidden = reinterpret_cast<at::Tensor*>(
+            mtp_pre_fc_norm_hidden_ptr);
+        auto* mtp_norm = reinterpret_cast<at::Tensor*>(mtp_norm_ptr);
+        TORCH_CHECK(!ctx->embed_ptr.empty() && ctx->embed_ptr[0] &&
+                !ctx->lm_head_ptr.empty() && ctx->lm_head_ptr[0] &&
+                ctx->embed_ptr[0]->dim() == 2 &&
+                ctx->lm_head_ptr[0]->dim() == 2,
+            "MTP requires valid vocabulary embedding and LM-head matrices");
+        const int64_t hidden_size = ctx->embed_ptr[0]->size(1);
+        const int64_t expected_vocab_rows = ctx->vocab_parallel
+            ? ctx->local_vocab_size : ctx->vocab_size;
+        TORCH_CHECK(ctx->embed_ptr[0]->size(0) == expected_vocab_rows &&
+                ctx->lm_head_ptr[0]->size(0) == expected_vocab_rows &&
+                ctx->lm_head_ptr[0]->size(1) == hidden_size,
+            "MTP requires embedding and LM-head shapes [local_vocab, hidden]");
+        const auto device = ctx->embed_ptr[0]->device();
+        for (const auto* tensor : {
+                mtp_fc, pre_norm_emb, pre_norm_hidden, mtp_norm}) {
+            TORCH_CHECK(tensor->defined() && tensor->is_cuda() &&
+                    tensor->device() == device &&
+                    tensor->scalar_type() == ctx->compute_type &&
+                    !tensor->requires_grad(),
+                "MTP tensors must be frozen CUDA tensors matching the base compute type");
+        }
+        TORCH_CHECK(mtp_fc->dim() == 2 &&
+                mtp_fc->sizes() == at::IntArrayRef({hidden_size, 2 * hidden_size}),
+            "MTP projection must have shape [hidden, 2 * hidden]");
+        for (const auto* tensor : {pre_norm_emb, pre_norm_hidden, mtp_norm}) {
+            TORCH_CHECK(tensor->dim() == 1 && tensor->size(0) == hidden_size,
+                "MTP normalization tensors must have shape [hidden]");
+        }
 
-        auto lm_head_tile = lm_head.narrow(0, v_start, v_n);
-        auto logits_tile = at::matmul(hidden_flat, lm_head_tile.t()).to(at::kFloat);
+        auto* layer_configs = reinterpret_cast<LayerConfig*>(
+            mtp_layer_configs_ptr);
+        std::vector<LayerConfig> configured_layers(
+            layer_configs, layer_configs + num_mtp_layers);
+        for (auto& config : configured_layers) {
+            config.nccl_comm = ctx->expert_parallel ? ctx->nccl_comm : nullptr;
+            config.nccl_stream = ctx->expert_parallel
+                ? ctx->nccl_stream : nullptr;
+        }
+        int64_t expected_weight_count = 0;
+        for (const auto& config : configured_layers)
+            expected_weight_count += weight_count_for_layer(config);
+        TORCH_CHECK(expected_weight_count == num_mtp_layer_weights &&
+                mtp_layer_weight_ptrs,
+            "MTP layer weight count does not match configured prediction layers");
+        auto** layer_weights = reinterpret_cast<at::Tensor**>(
+            mtp_layer_weight_ptrs);
+        std::vector<at::Tensor*> configured_weights;
+        configured_weights.reserve(num_mtp_layer_weights);
+        for (int64_t i = 0; i < num_mtp_layer_weights; ++i) {
+            auto* tensor = layer_weights[i];
+            TORCH_CHECK(tensor && tensor->defined() && tensor->is_cuda() &&
+                    tensor->device() == device &&
+                    tensor->scalar_type() == ctx->compute_type &&
+                    !tensor->requires_grad(),
+                "MTP layer weights must be frozen CUDA tensors matching the base compute type");
+            configured_weights.push_back(tensor);
+        }
+        int64_t weight_offset = 0;
+        for (int64_t layer = 0; layer < num_mtp_layers; ++layer) {
+            validate_mtp_prediction_layer_layout(
+                ctx, configured_layers[layer], layer_weights, weight_offset,
+                layer, hidden_size);
+            weight_offset += weight_count_for_layer(configured_layers[layer]);
+        }
 
-        // Gather target logits: subtract v_start to get local index
-        auto local_targets = (shifted_targets - v_start).clamp_min(0);
-        // Gather: logits_tile[i, local_targets[i]] for tokens in range
-        auto gathered = at::gather(logits_tile, /*dim=*/1, local_targets.reshape({-1, 1}));
-        // Only keep tokens that are actually in range
-        gathered = gathered * in_range.to(at::kFloat).reshape({-1, 1});
-        target_logit = at::max(target_logit, gathered);
+        auto configured_fused_qkv = ctx->fused_full_attention_qkv_weights;
+        auto configured_fused_fc1 = ctx->fused_mlp_fc1_weights;
+        configured_fused_qkv.resize(ctx->num_layers + num_mtp_layers);
+        configured_fused_fc1.resize(ctx->num_layers + num_mtp_layers);
+        {
+            at::NoGradGuard guard;
+            weight_offset = 0;
+            for (int64_t layer = 0; layer < num_mtp_layers; ++layer) {
+                const int64_t cache_index = ctx->num_layers + layer;
+                if (env_enabled("QWEN36_FUSED_QKV", false)) {
+                    configured_fused_qkv[cache_index] = at::cat({
+                        *layer_weights[weight_offset + 2],
+                        *layer_weights[weight_offset + 4],
+                        *layer_weights[weight_offset + 6]}, 0).contiguous();
+                }
+                if (env_enabled("QWEN36_FUSED_MLP_FC1", false)) {
+                    const int64_t mlp_start = 8;
+                    const int64_t gate_offset = configured_layers[layer].num_experts > 0
+                        ? mlp_start + 2 : mlp_start;
+                    configured_fused_fc1[cache_index] = at::cat({
+                        *layer_weights[weight_offset + gate_offset],
+                        *layer_weights[weight_offset + gate_offset + 1]},
+                        0).contiguous();
+                }
+                weight_offset += weight_count_for_layer(
+                    configured_layers[layer]);
+            }
+        }
 
-        c10::cuda::CUDACachingAllocator::emptyCache();
+        // Commit only after every pointer and layout has passed validation so a
+        // failed setup cannot leave a half-configured collective context.
+        ctx->mtp_fc = mtp_fc;
+        ctx->mtp_pre_fc_norm_emb = pre_norm_emb;
+        ctx->mtp_pre_fc_norm_hidden = pre_norm_hidden;
+        ctx->mtp_norm = mtp_norm;
+        ctx->mtp_layer_weights = std::move(configured_weights);
+        ctx->mtp_layer_configs = std::move(configured_layers);
+        ctx->fused_full_attention_qkv_weights =
+            std::move(configured_fused_qkv);
+        ctx->fused_mlp_fc1_weights = std::move(configured_fused_fc1);
+        ctx->has_mtp = true;
+
+        fprintf(stderr, "[q36_ctx] MTP set: %ld MTP layers, %ld MTP weight pointers\n",
+            (long)num_mtp_layers, (long)num_mtp_layer_weights);
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] set_mtp_weights FAILED: %s\n", e.what());
+        return -1;
     }
+}
 
-    // Loss per token: log(sum_exp) - target_logit  (= -log(softmax[target]))
-    auto log_sum_exp = at::log(sum_exp) + logit_max;  // log(sum(exp(x-max))) + max = logsumexp
-    auto per_token_loss = log_sum_exp - target_logit;  // [total_tokens, 1]
-    auto masked_loss = per_token_loss.squeeze(1) * mask_f;
-    auto total_count = mask_f.sum().clamp_min(1.0);
-    double loss_val = (masked_loss.sum().item<double>()) / total_count.item<double>();
+__attribute__((visibility("default"))) int32_t qwen36_set_pad_token_id(
+    void* ctx_ptr, int64_t pad_token_id
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "null training context");
+        TORCH_CHECK(pad_token_id >= 0,
+            "pad_token_id must be non-negative");
+        ctx->pad_token_id = pad_token_id;
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] set_pad_token_id FAILED: %s\n", e.what());
+        return -1;
+    }
+}
 
-    // ── Backward pass: compute grad_hidden_normed manually ──
-    // dL/dhidden_normed = (softmax - one_hot) / count * mask
-    // softmax = exp(logit - logit_max) / sum_exp
-    // We iterate tiles again, accumulate grad = softmax_tile @ lm_head_tile + target_grad
-    auto grad_hidden = at::zeros({total_tokens, hidden_dim},
-        at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device()));
-    auto grad_scale = mask_f / total_count;  // [total_tokens]
+// One training micro-step. Non-final micro-steps accumulate scaled leaf
+// gradients; only the final micro-step synchronizes and updates parameters.
+__attribute__((visibility("default"))) double qwen36_train_micro_step(
+    void* ctx_ptr,
+    void* input_ids_ptr,
+    void* target_mask_ptr,
+    void* attention_mask_ptr,
+    double gradient_scale,
+    int32_t apply_optimizer
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        GradientAccumulationFailureGuard accumulation_guard{ctx};
+        if (ctx && ctx->pp_world_size > 1) {
+            auto* input_ids = reinterpret_cast<at::Tensor*>(input_ids_ptr);
+            auto* target_mask = reinterpret_cast<at::Tensor*>(target_mask_ptr);
+            auto* attention_mask = attention_mask_ptr
+                ? reinterpret_cast<at::Tensor*>(attention_mask_ptr)
+                : nullptr;
+            TORCH_CHECK(input_ids && target_mask,
+                "pipeline train step requires input and target tensors");
+            const double loss = qwen36_pipeline_train_micro_step(
+                ctx, *input_ids, *target_mask, attention_mask,
+                gradient_scale, apply_optimizer);
+            accumulation_guard.disarmed = true;
+            return loss;
+        }
+        validate_native_execution_topology_collective(ctx);
+        TORCH_CHECK(!ctx->has_mtp || ctx->ep_world_size == 1,
+            "fixed-LoRA MTP with expert parallelism remains disabled until "
+            "its MTP/main denominator uses global source token counts; use "
+            "dynamic multi-LoRA for EP MTP");
+        auto* input_ids_tensor = reinterpret_cast<at::Tensor*>(input_ids_ptr);
+        auto* target_mask_tensor = reinterpret_cast<at::Tensor*>(target_mask_ptr);
+        auto* attention_mask_tensor = attention_mask_ptr
+            ? reinterpret_cast<at::Tensor*>(attention_mask_ptr)
+            : nullptr;
+        bool local_preflight_valid = input_ids_tensor && target_mask_tensor &&
+            input_ids_tensor->is_cuda() && target_mask_tensor->is_cuda() &&
+            input_ids_tensor->device() == target_mask_tensor->device() &&
+            input_ids_tensor->dim() == 2 && target_mask_tensor->dim() == 2 &&
+            input_ids_tensor->size(1) > 1 &&
+            (!ctx->has_mtp || input_ids_tensor->size(1) >= 3) &&
+            input_ids_tensor->sizes() == target_mask_tensor->sizes() &&
+            input_ids_tensor->scalar_type() == at::kLong &&
+            supported_mask_dtype(*target_mask_tensor) &&
+            ctx->adapters.empty() && gradient_scale > 0.0 &&
+            std::isfinite(gradient_scale) &&
+            (apply_optimizer == 0 || apply_optimizer == 1);
+        if (local_preflight_valid && attention_mask_tensor) {
+            local_preflight_valid = attention_mask_tensor->is_cuda() &&
+                attention_mask_tensor->device() == input_ids_tensor->device() &&
+                attention_mask_tensor->dim() == 2 &&
+                attention_mask_tensor->sizes() == input_ids_tensor->sizes() &&
+                supported_mask_dtype(*attention_mask_tensor);
+        }
+        if (local_preflight_valid &&
+            (ctx->nccl_comm || ctx->dp_comm || ctx->tp_comm ||
+                ctx->cp_comm)) {
+            local_preflight_valid =
+                input_ids_tensor->device().index() == ctx->cuda_device;
+        }
+        TORCH_CHECK(adapter_collective_all_succeeded(
+            ctx, local_preflight_valid) && local_preflight_valid,
+            "native Qwen fixed-LoRA call must use CUDA inputs on the NCCL "
+            "context device, a finite positive gradient scale, sequence "
+            "length >= 3 when MTP is enabled, and no dynamic adapters");
+        auto& input_ids = *input_ids_tensor;
+        const int input_device = input_ids.device().index();
+        if (!ctx->nccl_comm && !ctx->dp_comm && !ctx->tp_comm &&
+                !ctx->cp_comm)
+            ctx->cuda_device = input_device;
+        c10::cuda::set_device(ctx->cuda_device);
+        cudaSetDevice(ctx->cuda_device);
+        const bool replica_inputs_match = replica_input_signatures_match(
+            ctx, input_ids, *target_mask_tensor, attention_mask_tensor);
+        TORCH_CHECK(adapter_collective_all_succeeded(
+                ctx, replica_inputs_match) && replica_inputs_match,
+            "native Qwen fixed-LoRA input shape, dtype, or value differs "
+            "across TP, CP, or replicated EP ranks");
+        const double supervised_tokens = target_mask_tensor
+            ->narrow(1, 1, target_mask_tensor->size(1) - 1)
+            .to(at::kFloat).sum().item<double>();
+        const double micro_token_weight =
+            gradient_scale * supervised_tokens;
+        const double next_accumulated_token_weight =
+            ctx->accumulated_token_weight + micro_token_weight;
+        const bool local_token_weight_valid =
+            std::isfinite(supervised_tokens) && supervised_tokens >= 0.0 &&
+            std::isfinite(micro_token_weight) &&
+            std::isfinite(next_accumulated_token_weight);
+        TORCH_CHECK(adapter_collective_all_succeeded(
+                ctx, local_token_weight_valid) && local_token_weight_valid,
+            "native Qwen fixed-LoRA token weights must remain finite and "
+            "non-negative on every distributed rank");
+        // Fail collective-layout or optimizer-clock disagreement before
+        // forward/backward launches any work. The failure guard aborts any
+        // pending accumulation window without leaving queued graph work.
+        validate_fixed_collective_registry(ctx, apply_optimizer);
+        auto& target_mask = *target_mask_tensor;
+        if (attention_mask_ptr) {
+            auto& attention_mask = *reinterpret_cast<at::Tensor*>(attention_mask_ptr);
+            validate_linear_attention_mask(ctx, attention_mask);
+            ctx->attention_mask = attention_mask;
+            elide_trivial_attention_mask(ctx);
+        } else if (ctx->pad_token_id >= 0) {
+            ctx->attention_mask = derive_attention_mask(ctx, input_ids);
+            elide_trivial_attention_mask(ctx);
+        } else {
+            ctx->attention_mask = at::Tensor();
+            ctx->attention_lengths = at::Tensor();
+        }
 
-    for (int64_t t = 0; t < num_tiles; t++) {
-        int64_t v_start = t * tile_size;
-        int64_t v_end = std::min(v_start + tile_size, vocab_size);
-        int64_t v_n = v_end - v_start;
+        // Forward: checkpoint (default) or fused layer (QWEN36_FUSED_LAYER=1)
+        bool use_fused = env_enabled("QWEN36_FUSED_LAYER");
+        TORCH_CHECK(!use_fused,
+            "QWEN36_FUSED_LAYER is disabled until its custom backward preserves the layer graph");
+        begin_router_aux_loss(ctx, micro_token_weight);
+        auto hidden = use_fused
+            ? forward_full_fused(ctx, input_ids)
+            : ctx->use_checkpoint
+                ? forward_full_checkpoint(ctx, input_ids)
+                : forward_full(ctx, input_ids);
+        const double router_aux_loss = finish_router_aux_loss(ctx);
 
-        auto lm_head_tile = lm_head.narrow(0, v_start, v_n);  // [v_n, hidden]
-        auto logits_tile = at::matmul(hidden_flat, lm_head_tile.t()).to(at::kFloat);
+        // Debug: GPU memory after forward
+        {
+            size_t free, total;
+            cudaMemGetInfo(&free, &total);
+        }
 
-        // softmax_tile = exp(logits - max) / sum_exp  (reuse from forward)
-        auto softmax_tile = at::exp(logits_tile - logit_max) / sum_exp;  // [total_tokens, v_n]
+        // Main loss — compute_loss does chunked CE with immediate backward per chunk.
+        // This avoids accumulating 250 chunks of [512, vocab] logits in autograd graph.
+        auto loss_hidden = sequence_parallel_loss_gather(ctx, hidden);
+        auto main_loss = compute_loss(
+            ctx, loss_hidden, input_ids, target_mask, ctx->vocab_size);
+        double loss_val = main_loss.value.item<double>() + router_aux_loss;
+        auto total_hidden_grad = main_loss.hidden_grad;
+
+        // MTP must be differentiated before the main model backward. Its
+        // frozen head still contributes a hidden-state gradient to trainable
+        // main-layer LoRA parameters.
+        if (ctx->has_mtp && !env_enabled("QWEN36_DISABLE_MTP")) {
+            auto mtp_input = hidden.detach().set_requires_grad(true);
+            auto mtp_hidden = mtp_forward(ctx, mtp_input, input_ids);
+            auto mtp_loss = mtp_compute_loss(ctx, mtp_hidden, input_ids, target_mask);
+            mtp_loss.value.backward();
+            TORCH_CHECK(mtp_input.grad().defined(), "MTP did not produce a hidden gradient");
+            const double mtp_tokens = target_mask.narrow(1, 2,
+                target_mask.size(1) - 2).to(at::kFloat).sum().item<double>();
+            const double mtp_to_main_scale = supervised_tokens > 0.0
+                ? mtp_tokens / supervised_tokens : 0.0;
+            // MTP is averaged over its own response-token set. Convert it to
+            // the main-loss denominator before the common token-weighted
+            // gradient path, including uneven masks and DP replicas.
+            total_hidden_grad.add_(mtp_input.grad() * mtp_to_main_scale);
+            loss_val += mtp_loss.value.item<double>() * mtp_to_main_scale;
+        }
 
-        // Subtract one_hot for target tokens in this tile
-        auto in_range = (shifted_targets >= v_start) & (shifted_targets < v_end);
-        if (in_range.any().item<bool>()) {
-            auto local_targets = (shifted_targets - v_start).clamp_min(0);
-            // scatter -1 at target positions
-            auto one_hot = at::zeros({total_tokens, v_n},
-                at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device()));
-            one_hot.scatter_(1, local_targets.reshape({-1, 1}), 1.0);
-            one_hot = one_hot * in_range.to(at::kFloat).reshape({-1, 1});
-            softmax_tile = softmax_tile - one_hot;
+        const bool local_loss_finite =
+            std::isfinite(loss_val) && loss_val >= 0.0 &&
+            at::isfinite(total_hidden_grad).all().item<bool>();
+        TORCH_CHECK(adapter_collective_all_succeeded(
+                ctx, local_loss_finite) && local_loss_finite,
+            "native Qwen fixed-LoRA loss or hidden gradient became non-finite");
+
+        // compute_loss returns a local token mean. Convert it to a weighted
+        // numerator before backward; the FP32 window is divided by the global
+        // accumulated token weight exactly once at the optimizer boundary.
+        total_hidden_grad.mul_(micro_token_weight);
+
+        // Trigger exactly one main-model backward with the combined hidden
+        // gradient. Manual groups are the non-autograd checkpoint fallback;
+        // normal and sub-checkpoint paths use the real graph.
+        if (!ctx->group_inputs.empty() && !env_enabled("QWEN36_SUBCKPT")) {
+            manual_group_backward(
+                ctx, sequence_parallel_plain_split(ctx, total_hidden_grad));
+        } else {
+            loss_hidden.backward(total_hidden_grad);
         }
 
-        // grad_hidden += grad_scale * softmax_tile @ lm_head_tile
-        // softmax_tile: [total_tokens, v_n] (Float), lm_head_tile: [v_n, hidden] (BF16)
-        // → [total_tokens, hidden] (Float)
-        auto grad_tile = at::matmul(
-            (softmax_tile * grad_scale.reshape({-1, 1})),
-            lm_head_tile.to(at::kFloat)
-        );
-        grad_hidden.add_(grad_tile);
+        // Consume every BF16 leaf gradient immediately. Only the FP32 buffers
+        // survive the micro-step boundary.
+        const bool accumulation_was_active = ctx->accumulation_active;
+        harvest_gradient_accumulators(ctx);
+        if (micro_token_weight == 0.0)
+            ctx->accumulation_active = accumulation_was_active;
+        ctx->accumulated_token_weight = next_accumulated_token_weight;
+
+        if (!apply_optimizer) {
+            const bool local_gradients_finite =
+                fixed_lora_accumulators_are_finite(ctx);
+            TORCH_CHECK(adapter_collective_all_succeeded(
+                    ctx, local_gradients_finite) && local_gradients_finite,
+                "native Qwen fixed-LoRA accumulated gradient became non-finite");
+            // The forward graph has been consumed, but parameters remain
+            // live for the next micro-batch. Never reuse cached LoRA deltas
+            // whose autograd nodes were freed by this backward.
+            ctx->lora_cache_valid = false;
+            ctx->lora_batch_valid = false;
+            accumulation_guard.disarmed = true;
+            return loss_val;
+        }
 
-        c10::cuda::CUDACachingAllocator::emptyCache();
+        // Replicated DP gradients are synchronized before the local Adam
+        // update. EP keeps replicated gradients local because its forward
+        // routed activation already contains the cross-rank sum.
+        apply_fixed_lora_optimizer(
+            ctx, target_mask, ctx->accumulated_token_weight);
+        accumulation_guard.disarmed = true;
+        return loss_val;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] train_step FAILED: %s\n", e.what());
+        return -1.0;
     }
+}
 
-    // Set gradient on hidden_normed (leaf tensor).
-    // grad_hidden covers [batch, seq-1, hidden] (shifted tokens).
-    // hidden_normed is [batch, seq, hidden] — pad first token with zeros.
-    auto grad_reshaped = grad_hidden.to(hidden_normed.scalar_type())
-        .reshape({hidden_normed.size(0), seq_len - 1, hidden_dim});
-    auto grad_full = at::cat({
-        at::zeros({hidden_normed.size(0), 1, hidden_dim},
-            at::TensorOptions().dtype(hidden_normed.scalar_type()).device(hidden_normed.device())),
-        grad_reshaped
-    }, /*dim=*/1);  // [batch, seq, hidden]
-    hidden_normed.mutable_grad() = grad_full;
+// Backward-compatible complete optimizer step.
+__attribute__((visibility("default"))) double qwen36_train_step(
+    void* ctx_ptr,
+    void* input_ids_ptr,
+    void* target_mask_ptr,
+    void* attention_mask_ptr
+) {
+    return qwen36_train_micro_step(
+        ctx_ptr, input_ids_ptr, target_mask_ptr, attention_mask_ptr, 1.0, 1);
+}
 
-    // Backprop hidden_normed gradient to hidden via rms_norm recompute.
-    if (hidden_normed.grad().defined()) {
-        hidden.set_requires_grad(true);
-        auto hidden_normed_recompute = rms_norm(hidden, final_norm, ctx->rms_eps);
-        hidden_normed_recompute.backward(hidden_normed.grad());
-    }
+// Get LoRA A tensor pointer by index
+__attribute__((visibility("default"))) void* qwen36_get_lora_a(void* ctx_ptr, int64_t index) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (index < 0 || index >= (int64_t)ctx->lora_a.size()) return nullptr;
+    return &ctx->lora_a[index];
+}
 
-    c10::cuda::CUDACachingAllocator::emptyCache();
+// Get LoRA B tensor pointer by index
+__attribute__((visibility("default"))) void* qwen36_get_lora_b(void* ctx_ptr, int64_t index) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (index < 0 || index >= (int64_t)ctx->lora_b.size()) return nullptr;
+    return &ctx->lora_b[index];
+}
 
-    return at::tensor({loss_val},
-        at::TensorOptions().dtype(at::kFloat).device(hidden.device()));
+// Read-only diagnostic accessor used by native validation and runtime
+// observability. The returned tensor is owned by the training context.
+__attribute__((visibility("default"))) void* qwen36_get_lora_grad_accumulator(
+    void* ctx_ptr, int64_t index, int32_t is_b
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx || index < 0 || index >= (int64_t)ctx->lora_a.size()) return nullptr;
+    auto& accumulators = is_b ? ctx->grad_accum_b : ctx->grad_accum_a;
+    if (index >= (int64_t)accumulators.size() || !accumulators[index].defined())
+        return nullptr;
+    return &accumulators[index];
 }
 
-// ──────────────────────────────────────────────────────────────────────
-// Conditional emptyCache — only syncs when GPU memory is running low.
-// This avoids the ~2s/step of GPU idle time caused by 24 unconditional
-// emptyCache calls (10 fwd groups + 10 bwd groups + 4 CE chunks).
-// ──────────────────────────────────────────────────────────────────────
-static void maybe_emptyCache(at::Device device, int64_t threshold_gb = 20) {
-    size_t free_mem, total_mem;
-    cudaMemGetInfo(&free_mem, &total_mem);
-    if ((int64_t)free_mem < threshold_gb * 1024 * 1024 * 1024) {
-        c10::cuda::CUDACachingAllocator::emptyCache();
+// Explicitly abort an incomplete accumulation window. The operation is
+// idempotent and leaves parameters, Adam state, and optimizer clocks intact.
+__attribute__((visibility("default"))) int32_t qwen36_abort_gradient_accumulation(
+    void* ctx_ptr
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "null training context");
+        clear_gradient_accumulators(ctx);
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] abort_gradient_accumulation FAILED: %s\n", e.what());
+        return -1;
     }
 }
 
-// Cross-entropy loss with response-only masking — chunked with detach.
-// Detach hidden_normed so CE backward doesn't traverse main model graph.
-// Accumulate hidden_normed gradient, then backprop to hidden separately.
-static at::Tensor compute_loss(
-    TrainingContext* ctx,
-    const at::Tensor& hidden,
-    const at::Tensor& input_ids,
-    const at::Tensor& target_mask,
-    int64_t vocab_size
+// Copy one exported LoRA tensor back into the native leaf parameter. This is
+// used by checkpoint resume and adapter import; derived delta caches are
+// invalidated so the next forward rebuilds the graph from the new leaf.
+__attribute__((visibility("default"))) int32_t qwen36_set_lora_tensor(
+    void* ctx_ptr, int64_t index, int32_t is_b, void* tensor_ptr
 ) {
-    auto final_norm = *ctx->final_norm_ptr[0];
-    auto lm_head = *ctx->lm_head_ptr[0];
-
-    // Detach hidden for CE computation — CE backward won't touch main model graph.
-    // We accumulate gradient into hidden_normed, then manually backprop to hidden.
-    auto hidden_detached = hidden.detach();
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(tensor_ptr, "null LoRA tensor");
+        auto& slots = is_b ? ctx->lora_b : ctx->lora_a;
+        TORCH_CHECK(index >= 0 && index < (int64_t)slots.size(), "invalid LoRA slot");
+        auto& target = slots[index];
+        auto& source = *reinterpret_cast<at::Tensor*>(tensor_ptr);
+        TORCH_CHECK(source.sizes() == target.sizes(),
+            "LoRA tensor shape mismatch at slot ", index,
+            ": expected ", target.sizes(), " got ", source.sizes());
+        at::NoGradGuard guard;
+        target.copy_(source.to(target.device()).to(target.scalar_type()));
+        ctx->lora_cache_valid = false;
+        ctx->lora_batch_valid = false;
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] set_lora_tensor FAILED: %s\n", e.what());
+        return -1;
+    }
+}
 
-    // Compute hidden_normed in no-grad, then set requires_grad on it.
-    // This way CE backward only builds a tiny graph (hidden_normed → logits → loss),
-    // not connected to hidden_detached at all.
-    at::Tensor hidden_normed;
-    {
-        at::AutoGradMode no_grad_mode(false);
-        hidden_normed = rms_norm(hidden_detached, final_norm, ctx->rms_eps);
+// Free training context
+__attribute__((visibility("default"))) void qwen36_free_training_context(void* ctx_ptr) {
+    if (ctx_ptr) {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        // Don't destroy NCCL communicator — it's a process-level singleton
+        // (g_nccl_comm). It must survive context destruction so the next
+        // session can reuse it. Destroying it would break NCCL for all
+        // subsequent sessions in the same worker process.
+        // ncclCommDestroy is called only on process exit (via atexit or Drop).
+        if (ctx->moe_shared_stream || ctx->moe_ep_metadata_stream) {
+            int previous_device = 0;
+            if (cudaGetDevice(&previous_device) == cudaSuccess) {
+                if (previous_device != ctx->cuda_device)
+                    cudaSetDevice(ctx->cuda_device);
+                if (ctx->moe_shared_stream)
+                    cudaStreamDestroy(ctx->moe_shared_stream);
+                if (ctx->moe_ep_metadata_stream)
+                    cudaStreamDestroy(ctx->moe_ep_metadata_stream);
+                if (previous_device != ctx->cuda_device)
+                    cudaSetDevice(previous_device);
+            }
+            ctx->moe_shared_stream = nullptr;
+            ctx->moe_ep_metadata_stream = nullptr;
+        }
+        delete ctx;
     }
-    hidden_normed.set_requires_grad(true);
+}
 
-    int64_t seq_len = hidden_normed.size(1);
-    auto shifted_hidden = hidden_normed.narrow(1, 0, seq_len - 1);
-    auto shifted_targets = input_ids.narrow(1, 1, seq_len - 1).reshape({-1});
-    auto shifted_mask = target_mask.narrow(1, 1, seq_len - 1).reshape({-1});
+// ── Batched Multi-LoRA Training ──
 
-    int64_t total_tokens = shifted_targets.size(0);
-    // Smaller chunks = less peak memory (4GB vs 16GB per chunk at vocab=248K).
-    // This allows removing emptyCache between chunks — GPU stays async.
-    int64_t chunk_size = 4096;
-    int64_t num_chunks = (total_tokens + chunk_size - 1) / chunk_size;
+/// Compute max adapters that fit in available GPU memory.
+/// Based on per-adapter activation memory (dominant) + LoRA params + Adam state.
+static int64_t compute_n_max(
+    int64_t free_gpu_bytes, int64_t rank, int64_t seq,
+    int64_t hidden, int64_t group_size, int64_t num_layers,
+    bool vocab_parallel, int64_t local_vocab_size
+) {
+    // Per-adapter activation memory (BF16, group_size=2 optimal):
+    // group_inputs: (num_layers / group_size) × seq × hidden × 2 bytes
+    // peak recompute: ~34 MB per layer pair × group_size (rough avg)
+    int64_t gs = group_size < 1 ? 1 : group_size;
+    int64_t num_groups = (num_layers + gs - 1) / gs;
+    int64_t group_input_mem = num_groups * seq * hidden * 2;  // BF16
+
+    // Average per-layer saved tensors during recompute:
+    // full_attn: ~17MB, linear_attn: ~25MB, moe: ~8MB → avg ~17MB
+    int64_t avg_layer_saved = 17 * 1024 * 1024;  // bytes
+    int64_t peak_mem = gs * avg_layer_saved;
+
+    // LoRA params + Adam state: 280 modules × (A+B) × (BF16 param + FP32 m + FP32 v)
+    // Per module: rank × hidden × 2 (BF16) + hidden × rank × 2 (BF16) + 2 × 4 (FP32 m+v)
+    // Simplified: 280 × rank × hidden × (2 + 2 + 8) = 280 × rank × hidden × 12
+    int64_t num_modules = 280;
+    int64_t lora_mem = num_modules * rank * hidden * 12;  // conservative
+
+    // Vocabulary TP keeps a cached FP32 logits tile while backward also owns
+    // FP32 softmax and one-hot buffers. Reserve four tile-sized buffers plus
+    // the [token, hidden] gradient so tenant chunking stays conservative.
+    int64_t ce_peak = vocab_parallel
+        ? 512LL * local_vocab_size * 16LL + 512LL * hidden * 4LL
+        : 16384LL * 248320LL * 4LL;
+
+    // Attention intermediate: Q/K/V + attn_weights ≈ 4 × N × heads × seq × head_dim × 2 bytes
+    // Flash attention keeps this O(seq) not O(seq²), but still significant at seq=16K
+    int64_t attn_heads = 16;
+    int64_t head_dim = 256;
+    int64_t attn_mem = 4 * attn_heads * seq * head_dim * 2;  // per adapter, BF16
+
+    int64_t per_adapter = group_input_mem + peak_mem + lora_mem + attn_mem;
+    // Empirical multiplier: residual add, MoE routing, LoRA delta, etc.
+    // Multiplier scales with seq: small seq needs less (CPU dispatch dominant),
+    // large seq needs more (attention/MoE intermediates dominate).
+    // seq=512: 3x → n_max=100. seq=16K: 8x → n_max≈8 (auto-chunks N=20+).
+    int64_t mult = (seq > 4096) ? 8 : 3;
+    per_adapter = per_adapter * mult;
+    if (per_adapter <= 0) return 1;
+
+    // Reserve 15% for fragmentation + overhead, minus CE peak (constant overhead)
+    int64_t usable = (free_gpu_bytes - ce_peak) * 85 / 100;
+    if (usable < per_adapter) usable = free_gpu_bytes * 50 / 100;  // fallback: aggressive
+    int64_t n_max = usable / per_adapter;
+    return n_max < 1 ? 1 : n_max;
+}
+
+enum class DynamicMultiLoraMode {
+    TrainAndFinalize,
+    TrainOnly,
+    FinalizeOnly,
+};
 
-    auto total_count = shifted_mask.sum().clamp_min(1.0);
-    auto hidden_flat = shifted_hidden.reshape({-1, hidden_normed.size(2)});
+/// Train all adapters in chunks. Inputs may be [1, seq] (shared prompt,
+/// repeated per chunk) or [n_total, seq] (one independent sample per adapter).
+/// Heterogeneous selected training reuses the same implementation in two
+/// phases so every selected adapter shares one synchronization/Adam boundary.
+double qwen36_train_multi_lora_impl(
+    void* ctx_ptr,
+    void* input_ids_ptr,
+    void* target_mask_ptr,
+    void* attention_mask_ptr,
+    int32_t n_total,
+    int32_t lora_rank,
+    DynamicMultiLoraMode mode,
+    int32_t* finalizer_phase = nullptr,
+    at::Tensor* loss_numerators_out = nullptr,
+    at::Tensor* token_counts_out = nullptr,
+    bool allow_pipeline = false,
+    const at::Tensor* token_counts_override = nullptr,
+    bool accumulators_are_numerators = false,
+    bool retain_adam_shadows = false
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        DynamicAdamShadowLeaseScope shadow_scope{
+            ctx, retain_adam_shadows, false};
+        const bool train_only = mode == DynamicMultiLoraMode::TrainOnly;
+        const bool finalize_only = mode == DynamicMultiLoraMode::FinalizeOnly;
+        if (finalize_only && finalizer_phase) *finalizer_phase = 0;
+        if (mode == DynamicMultiLoraMode::TrainAndFinalize)
+            validate_dynamic_context_health(ctx);
+        validate_native_execution_topology_collective(ctx, allow_pipeline);
+        auto* input_ids_tensor = reinterpret_cast<at::Tensor*>(input_ids_ptr);
+        auto* target_mask_tensor = reinterpret_cast<at::Tensor*>(target_mask_ptr);
+        auto* attention_mask_tensor = attention_mask_ptr
+            ? reinterpret_cast<at::Tensor*>(attention_mask_ptr)
+            : nullptr;
+        const int64_t total_adapters = (int64_t)ctx->adapters.size();
+        const bool mtp_enabled =
+            ctx->has_mtp && !env_enabled("QWEN36_DISABLE_MTP");
+        const bool dynamic_mtp_tp_supported = ctx->tp_world_size == 1 ||
+            (ctx->tp_world_size == 2 && ctx->tp_comm &&
+                !ctx->sequence_parallel &&
+                ctx->base_tp_attention && ctx->base_tp_mlp);
+        const bool dynamic_mtp_ep_supported = ctx->ep_world_size == 1 ||
+            (ctx->expert_parallel && ctx->nccl_comm &&
+                env_enabled("QWEN36_EP_A2A") &&
+                env_enabled("QWEN36_EP_A2A_SHARDED") &&
+                env_enabled("QWEN36_EP_A2A_PACKED", true));
+        const bool dynamic_mtp_topology_supported = !mtp_enabled ||
+            (!allow_pipeline && dynamic_mtp_tp_supported &&
+                dynamic_mtp_ep_supported && ctx->cp_world_size == 1 &&
+                ctx->pp_world_size == 1 &&
+                ctx->router_aux_loss_coef == 0.0);
+        bool local_input_valid = input_ids_tensor && target_mask_tensor &&
+            input_ids_tensor->is_cuda() && target_mask_tensor->is_cuda() &&
+            input_ids_tensor->device() == target_mask_tensor->device() &&
+            input_ids_tensor->dim() == 2 && target_mask_tensor->dim() == 2 &&
+            input_ids_tensor->size(1) > 1 &&
+            input_ids_tensor->scalar_type() == at::kLong &&
+            supported_mask_dtype(*target_mask_tensor) &&
+            target_mask_tensor->sizes() == input_ids_tensor->sizes() &&
+            n_total > 0 &&
+            (input_ids_tensor->size(0) == 1 ||
+                input_ids_tensor->size(0) == n_total) &&
+            (mode != DynamicMultiLoraMode::TrainAndFinalize ||
+                (!ctx->accumulation_active &&
+                    ctx->accumulated_token_weight == 0.0)) &&
+            dynamic_mtp_topology_supported &&
+            (!mtp_enabled || input_ids_tensor->size(1) >= 3) &&
+            (!mtp_enabled || !token_counts_override);
+        if (local_input_valid && attention_mask_tensor) {
+            local_input_valid = attention_mask_tensor->is_cuda() &&
+                attention_mask_tensor->device() == input_ids_tensor->device() &&
+                attention_mask_tensor->dim() == 2 &&
+                attention_mask_tensor->sizes() == input_ids_tensor->sizes() &&
+                supported_mask_dtype(*attention_mask_tensor);
+        }
+        if (local_input_valid &&
+            (ctx->nccl_comm || ctx->dp_comm || ctx->tp_comm ||
+                ctx->cp_comm)) {
+            local_input_valid =
+                input_ids_tensor->device().index() == ctx->cuda_device;
+        }
+        TORCH_CHECK(adapter_collective_all_succeeded(
+                ctx, local_input_valid) && local_input_valid,
+            "native Qwen multi-LoRA inputs must be CUDA tensors on the "
+            "NCCL context device");
+        TORCH_CHECK(dynamic_mtp_topology_supported,
+            "dynamic MTP currently requires CP=PP=1, DP>=1, compatible "
+            "vocabulary and replicated sequence state, sharded prediction-layer weights "
+            "when TP>1, packed sharded A2A when EP>1, and router aux disabled");
+        auto& input_ids = *input_ids_tensor;
+        auto& target_mask = *target_mask_tensor;
+        const int input_device = input_ids.device().index();
+        if (!ctx->nccl_comm && !ctx->dp_comm && !ctx->tp_comm &&
+                !ctx->cp_comm)
+            ctx->cuda_device = input_device;
+        c10::cuda::set_device(ctx->cuda_device);
+        cudaSetDevice(ctx->cuda_device);
+        const bool replica_inputs_match = replica_input_signatures_match(
+            ctx, input_ids, target_mask, attention_mask_tensor);
+        TORCH_CHECK(adapter_collective_all_succeeded(
+                ctx, replica_inputs_match) && replica_inputs_match,
+            "native Qwen multi-LoRA input shape and dtype differ across TP "
+            "or replicated EP ranks");
+        GradientAccumulationFailureGuard accumulation_guard{ctx};
+
+        validate_adapter_collective_registry(
+            ctx, nullptr, n_total, finalize_only ? 0 : lora_rank,
+            true);
+        if (total_adapters == 0) return -1.0;
+        TORCH_CHECK(n_total > 0 && total_adapters == n_total,
+            "n_total must equal the number of registered adapters (n_total=",
+            n_total, ", registered=", total_adapters, ")");
+        if (!finalize_only) {
+            const auto& reference_adapter = ctx->adapters.front();
+            int64_t maximum_rank = 0;
+            for (const auto& adapter : ctx->adapters) {
+                maximum_rank = std::max(maximum_rank, adapter.rank);
+                if (!ctx->pad_heterogeneous_lora_batch) {
+                    TORCH_CHECK(adapter.rank == lora_rank,
+                        "lora_rank argument must match every registered adapter; adapter=",
+                        adapter.id, " registered_rank=", adapter.rank,
+                        " requested_rank=", lora_rank);
+                    TORCH_CHECK(
+                        adapter.all_target_layers ==
+                            reference_adapter.all_target_layers &&
+                            adapter.global_target_layers ==
+                                reference_adapter.global_target_layers &&
+                            adapter.target_modules == reference_adapter.target_modules,
+                        "legacy multi-LoRA training requires homogeneous target layers/modules; "
+                        "use the v2 trainer for heterogeneous adapters");
+                }
+            }
+            TORCH_CHECK(!ctx->pad_heterogeneous_lora_batch ||
+                    maximum_rank == lora_rank,
+                "heterogeneous multi-LoRA rank hint must equal the maximum "
+                "registered rank; maximum=", maximum_rank,
+                " requested=", lora_rank);
+            ++ctx->multi_lora_invocation;
+        }
+        const int64_t input_batch = input_ids.size(0);
+        auto input_row_token_counts = target_mask
+            .narrow(1, 1, target_mask.size(1) - 1)
+            .to(at::kFloat).sum(1);
+        auto adapter_token_counts = input_batch == 1
+            ? input_row_token_counts.repeat({total_adapters})
+            : input_row_token_counts;
+        if (token_counts_override) {
+            TORCH_CHECK(token_counts_override->defined() &&
+                    token_counts_override->dim() == 1 &&
+                    token_counts_override->size(0) == total_adapters,
+                "dynamic token-count override must match adapter registry");
+            adapter_token_counts = token_counts_override->to(at::kFloat)
+                .to(input_ids.device()).contiguous();
+        }
+        at::Tensor mtp_adapter_token_counts;
+        at::Tensor mtp_gradient_scales;
+        at::Tensor mtp_report_scales;
+        if (mtp_enabled && !finalize_only) {
+            auto input_row_mtp_counts = target_mask
+                .narrow(1, 2, target_mask.size(1) - 2)
+                .to(at::kFloat).sum(1);
+            mtp_adapter_token_counts = input_batch == 1
+                ? input_row_mtp_counts.repeat({total_adapters})
+                : input_row_mtp_counts;
+            bool local_counts_valid = false;
+            try {
+                local_counts_valid = at::logical_and(
+                        at::isfinite(mtp_adapter_token_counts),
+                        at::logical_and(
+                            mtp_adapter_token_counts >= 0,
+                            mtp_adapter_token_counts <= adapter_token_counts))
+                    .all().item<bool>();
+            } catch (...) {
+                local_counts_valid = false;
+            }
+            TORCH_CHECK(adapter_collective_all_succeeded(
+                    ctx, local_counts_valid) && local_counts_valid,
+                "dynamic MTP token counts must be finite, non-negative, and no "
+                "larger than the corresponding main-loss token counts");
+
+            auto packed_counts = at::cat({
+                adapter_token_counts.to(at::kFloat),
+                mtp_adapter_token_counts}, 0).contiguous();
+            TORCH_CHECK(replica_token_weights_match(
+                    ctx, packed_counts,
+                    reinterpret_cast<ncclComm_t>(ctx->tp_comm),
+                    ctx->tp_world_size, "TP MTP"),
+                "dynamic MTP token counts differ across TP ranks");
+            const bool sharded_ep = ctx->expert_parallel && ctx->nccl_comm &&
+                env_enabled("QWEN36_EP_A2A_SHARDED");
+            if (ctx->expert_parallel && ctx->nccl_comm && !sharded_ep) {
+                TORCH_CHECK(replica_token_weights_match(
+                        ctx, packed_counts,
+                        reinterpret_cast<ncclComm_t>(ctx->nccl_comm),
+                        ctx->ep_world_size, "replicated EP MTP"),
+                    "dynamic MTP token counts differ across replicated EP ranks");
+            }
+            auto sum_counts = [&](ncclComm_t communicator, const char* axis) {
+                auto reduced_counts = at::empty_like(packed_counts);
+                auto stream = c10::cuda::getCurrentCUDAStream(
+                    packed_counts.device().index()).stream();
+                auto error = ncclAllReduce(
+                    packed_counts.data_ptr<float>(),
+                    reduced_counts.data_ptr<float>(), packed_counts.numel(),
+                    ncclFloat, ncclSum, communicator, stream);
+                TORCH_CHECK(error == ncclSuccess,
+                    "dynamic MTP token-count ", axis, " reduction failed: ",
+                    ncclGetErrorString(error));
+                packed_counts = reduced_counts;
+            };
+            if (sharded_ep) {
+                sum_counts(
+                    reinterpret_cast<ncclComm_t>(ctx->nccl_comm),
+                    "sharded EP");
+            }
+            if (ctx->data_parallel && ctx->dp_comm &&
+                    ctx->dp_world_size > 1) {
+                sum_counts(
+                    reinterpret_cast<ncclComm_t>(ctx->dp_comm), "DP");
+            }
+            auto global_main_counts = packed_counts.narrow(
+                0, 0, total_adapters);
+            auto global_mtp_counts = packed_counts.narrow(
+                0, total_adapters, total_adapters);
+            // MTP sample numerators already carry the MTP loss coefficient
+            // and are token sums.  The report denominator is the main-loss
+            // token count, so no extra MTP/main ratio belongs here; dividing
+            // the sum by global_main_counts applies that ratio exactly once.
+            mtp_report_scales = at::ones_like(global_main_counts);
+            mtp_gradient_scales = at::where(
+                at::logical_and(
+                    adapter_token_counts > 0, global_mtp_counts > 0),
+                (mtp_adapter_token_counts /
+                    adapter_token_counts.clamp_min(1.0)) *
+                    (global_main_counts / global_mtp_counts.clamp_min(1.0)),
+                at::zeros_like(adapter_token_counts));
+        }
+        // Keep the caller's mask intact. Each chunk receives either the
+        // corresponding rows or a repeated batch-1 mask; this also prevents
+        // elide_trivial_attention_mask from leaking the last chunk into the
+        // next train step.
+        at::Tensor provided_attention_mask;
+        if (attention_mask_tensor) {
+            provided_attention_mask = *attention_mask_tensor;
+            validate_linear_attention_mask(ctx, provided_attention_mask);
+        } else if (ctx->pad_token_id >= 0) {
+            provided_attention_mask = derive_attention_mask(ctx, input_ids);
+            validate_linear_attention_mask(ctx, provided_attention_mask);
+        }
+        const at::Tensor saved_attention_mask = ctx->attention_mask;
+        const at::Tensor saved_attention_lengths = ctx->attention_lengths;
+        const bool saved_use_checkpoint = ctx->use_checkpoint;
+        struct AttentionMaskGuard {
+            TrainingContext* ctx;
+            at::Tensor saved;
+            at::Tensor saved_lengths;
+            ~AttentionMaskGuard() {
+                ctx->attention_mask = saved;
+                ctx->attention_lengths = saved_lengths;
+            }
+        } attention_mask_guard{
+            ctx, saved_attention_mask, saved_attention_lengths};
+        if (!provided_attention_mask.defined()) {
+            ctx->attention_mask = at::Tensor();
+            ctx->attention_lengths = at::Tensor();
+        }
+        struct CheckpointModeGuard {
+            TrainingContext* ctx;
+            bool saved;
+            ~CheckpointModeGuard() { ctx->use_checkpoint = saved; }
+        } checkpoint_mode_guard{ctx, saved_use_checkpoint};
+
+        struct AdapterRegistryChunkGuard {
+            TrainingContext* ctx;
+            std::vector<TrainingContext::LoRAAdapter> all;
+            int64_t start = 0;
+            size_t moved = 0;
+            bool active = false;
+            bool adapter_id_index_was_valid = false;
+
+            AdapterRegistryChunkGuard(
+                TrainingContext* context, int64_t start, int64_t end,
+                bool scope_registry)
+                : ctx(context), start(start),
+                  adapter_id_index_was_valid(
+                      context && context->adapter_id_index_valid) {
+                if (!scope_registry) return;
+                TORCH_CHECK(start >= 0 && end >= start &&
+                        end <= static_cast<int64_t>(ctx->adapters.size()),
+                    "invalid dynamic adapter chunk range [", start, ", ",
+                    end, ") for registry size ", ctx->adapters.size());
+                ctx->adapter_id_index_valid = false;
+                all.swap(ctx->adapters);
+                try {
+                    const size_t count = static_cast<size_t>(end - start);
+                    ctx->adapters.reserve(count);
+                    for (int64_t index = start; index < end; ++index) {
+                        ctx->adapters.emplace_back(std::move(all[index]));
+                        ++moved;
+                    }
+                    active = true;
+                } catch (...) {
+                    for (size_t index = 0; index < moved; ++index)
+                        all[static_cast<size_t>(start) + index] =
+                            std::move(ctx->adapters[index]);
+                    ctx->adapters.clear();
+                    ctx->adapters.swap(all);
+                    ctx->adapter_id_index_valid =
+                        adapter_id_index_was_valid;
+                    throw;
+                }
+            }
 
-    double total_loss_val = 0.0;
-    auto total_count_val = total_count.item<double>();
+            void restore() {
+                if (!active) return;
+                for (size_t index = 0; index < moved; ++index)
+                    all[static_cast<size_t>(start) + index] =
+                        std::move(ctx->adapters[index]);
+                ctx->adapters.clear();
+                ctx->adapters.swap(all);
+                ctx->adapter_id_index_valid = adapter_id_index_was_valid;
+                moved = 0;
+                active = false;
+            }
 
-    for (int64_t c = 0; c < num_chunks; c++) {
-        int64_t start = c * chunk_size;
-        int64_t end = std::min(start + chunk_size, total_tokens);
-        int64_t n = end - start;
+            ~AdapterRegistryChunkGuard() { restore(); }
+        };
 
-        auto chunk_hidden = hidden_flat.narrow(0, start, n);
-        auto chunk_logits = at::matmul(chunk_hidden, lm_head.t());
-        auto chunk_targets = shifted_targets.narrow(0, start, n);
-        auto chunk_mask = shifted_mask.narrow(0, start, n);
+        int64_t n_max = total_adapters;
+        if (!finalize_only) {
+            // All workers in the TP x CP x EP x DP grid must agree on the
+            // activation chunk schedule to keep forward collectives ordered.
+            size_t free_mem = 0;
+            size_t total_mem = 0;
+            cudaMemGetInfo(&free_mem, &total_mem);
+            n_max = compute_n_max(
+                (int64_t)free_mem, lora_rank,
+                input_ids.size(-1), 2048,
+                ctx->group_size, ctx->num_layers, ctx->vocab_parallel,
+                ctx->local_vocab_size
+            );
+            n_max = std::min(n_max, total_adapters);
+            if (n_max < 1) n_max = 1;
+            if ((ctx->nccl_comm && ctx->ep_world_size > 1) ||
+                (ctx->dp_comm && ctx->dp_world_size > 1) ||
+                (ctx->tp_comm && ctx->tp_world_size > 1) ||
+                (ctx->cp_comm && ctx->cp_world_size > 1)) {
+                auto published_n_max = at::full(
+                    {1}, n_max, input_ids.options().dtype(at::kLong));
+                auto stream = c10::cuda::getCurrentCUDAStream(
+                    input_ids.device().index()).stream();
+                if (ctx->cp_comm && ctx->cp_world_size > 1) {
+                    auto err = ncclAllReduce(
+                        published_n_max.data_ptr<int64_t>(),
+                        published_n_max.data_ptr<int64_t>(), 1, ncclInt64,
+                        ncclMin, reinterpret_cast<ncclComm_t>(ctx->cp_comm),
+                        stream);
+                    TORCH_CHECK(err == ncclSuccess,
+                        "CP n_max all-reduce failed: ", ncclGetErrorString(err));
+                }
+                if (ctx->nccl_comm && ctx->ep_world_size > 1) {
+                    auto err = ncclAllReduce(
+                        published_n_max.data_ptr<int64_t>(),
+                        published_n_max.data_ptr<int64_t>(), 1, ncclInt64,
+                        ncclMin, reinterpret_cast<ncclComm_t>(ctx->nccl_comm),
+                        stream);
+                    TORCH_CHECK(err == ncclSuccess,
+                        "EP n_max all-reduce failed: ", ncclGetErrorString(err));
+                }
+                if (ctx->dp_comm && ctx->dp_world_size > 1) {
+                    auto err = ncclAllReduce(
+                        published_n_max.data_ptr<int64_t>(),
+                        published_n_max.data_ptr<int64_t>(), 1, ncclInt64,
+                        ncclMin, reinterpret_cast<ncclComm_t>(ctx->dp_comm),
+                        stream);
+                    TORCH_CHECK(err == ncclSuccess,
+                        "DP n_max all-reduce failed: ", ncclGetErrorString(err));
+                }
+                if (ctx->tp_comm && ctx->tp_world_size > 1) {
+                    auto err = ncclAllReduce(
+                        published_n_max.data_ptr<int64_t>(),
+                        published_n_max.data_ptr<int64_t>(), 1, ncclInt64,
+                        ncclMin, reinterpret_cast<ncclComm_t>(ctx->tp_comm),
+                        stream);
+                    TORCH_CHECK(err == ncclSuccess,
+                        "TP n_max all-reduce failed: ", ncclGetErrorString(err));
+                }
+                n_max = published_n_max.to(
+                    at::TensorOptions().device(at::kCPU)).item<int64_t>();
+            }
+            n_max = std::min(n_max, total_adapters);
+            if (n_max < 1) n_max = 1;
+            if (env_enabled("QWEN36_TRAIN_TRACE")) {
+                fprintf(stderr,
+                    "[train_multi] total=%ld n_max=%ld free=%.1fGB rank=%d\n",
+                    (long)total_adapters, (long)n_max,
+                    (double)free_mem / 1e9, lora_rank);
+            }
+        }
 
-        // Diagnostic: print logits stats for first chunk, first token
-        if (c == 0 && getenv("QWEN36_LOSS_DIAG")) {
-            auto logits_f = chunk_logits[0].to(at::kFloat);
-            auto lsm = at::log_softmax(logits_f, -1);
-            int64_t tgt = chunk_targets[0].item<int64_t>();
-            fprintf(stderr, "[diag] logits shape: [%ld, %ld]\n", (long)n, (long)chunk_logits.size(1));
-            fprintf(stderr, "[diag] logits[0,:5]: %.6f %.6f %.6f %.6f %.6f\n",
-                    logits_f[0].item<float>(), logits_f[1].item<float>(),
-                    logits_f[2].item<float>(), logits_f[3].item<float>(),
-                    logits_f[4].item<float>());
-            fprintf(stderr, "[diag] logits mean: %.6f, std: %.6f\n",
-                    logits_f.mean().item<float>(), logits_f.std().item<float>());
-            fprintf(stderr, "[diag] target token: %ld\n", (long)tgt);
-            fprintf(stderr, "[diag] target logit: %.6f\n", logits_f[tgt].item<float>());
-            fprintf(stderr, "[diag] log_softmax[target]: %.6f\n", lsm[tgt].item<float>());
-            fprintf(stderr, "[diag] -log_softmax[target] (per-token loss): %.6f\n", -lsm[tgt].item<float>());
+        at::Tensor total_loss;
+        at::Tensor step_numerics_valid;
+        if (!finalize_only) {
+            total_loss = at::zeros({}, input_ids.options().dtype(at::kDouble));
+            step_numerics_valid = at::ones(
+                {}, input_ids.options().dtype(at::kBool));
         }
+        at::Tensor loss_numerators;
+        at::Tensor loss_token_counts;
+        if (!finalize_only && loss_numerators_out && token_counts_out) {
+            loss_numerators = at::zeros({total_adapters},
+                at::TensorOptions().dtype(at::kFloat).device(input_ids.device()));
+            loss_token_counts = at::zeros_like(loss_numerators);
+        }
+        int64_t num_chunks = (total_adapters + n_max - 1) / n_max;
+        // Chunking is a memory scheduling detail, not an optimizer step. Each
+        // tenant clock commits only after its Adam launch succeeds.
 
-        auto per_token_loss = at::cross_entropy_loss(
-            chunk_logits.to(at::kFloat), chunk_targets,
-            at::Tensor(), at::Reduction::None, -100, 0.0
-        );
-        auto masked_loss = per_token_loss * chunk_mask.to(at::kFloat);
-        auto chunk_loss = masked_loss.sum();
+        for (int64_t chunk = 0; chunk < num_chunks; chunk++) {
+            int64_t start = chunk * n_max;
+            int64_t end = std::min(start + n_max, total_adapters);
+            int64_t n = end - start;
 
-        // Backward this chunk — each chunk creates an independent CE subgraph
-        // because hidden_normed is a leaf tensor. retain_graph=false is safe
-        // and much faster than retain_graph=true (which accumulates graph).
-        torch::autograd::backward({chunk_loss}, {},
-            /*retain_graph=*/false, /*create_graph=*/false);
+            // Invalidate cache for this chunk's adapter set
+            ctx->lora_batch_valid = false;
+            ctx->lora_cache_valid = false;
 
-        total_loss_val += chunk_loss.item<double>();
+            // Scope the registry to this activation-memory chunk. The guard
+            // restores the full registry even if forward/backward fails.
+            AdapterRegistryChunkGuard registry_guard(
+                ctx, start, end, !finalize_only);
+
+            at::Tensor chunk_loss;
+            if (!finalize_only) {
+                // Mark batched mode active
+                ctx->lora_batch_valid = true;  // triggers prepare_lora_batch in forward
+
+            // Expand only a shared batch-1 sample. With a tenant-specific
+            // [n_total, seq] batch, preserve each adapter's own row and slice
+            // the same chunk range as the adapter registry.
+            auto ids_expanded = input_batch == 1
+                ? input_ids.repeat({n, 1})
+                : input_ids.narrow(0, start, n).contiguous();
+            auto mask_expanded = input_batch == 1
+                ? target_mask.repeat({n, 1})
+                : target_mask.narrow(0, start, n).contiguous();
+            if (provided_attention_mask.defined()) {
+                ctx->attention_mask = input_batch == 1
+                    ? provided_attention_mask.repeat({n, 1})
+                    : provided_attention_mask.narrow(0, start, n).contiguous();
+                elide_trivial_attention_mask(ctx);
+            }
 
-        if (hidden_normed.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache();
-    }
+            // Run train_step (reuses existing forward + loss + backward + Adam)
+            // But we need to pass the expanded tensors
+            auto& input_ref = ids_expanded;
+            auto& mask_ref = mask_expanded;
 
-    // Backprop hidden_normed gradient to hidden via rms_norm.
-    // hidden_normed was computed in no-grad mode (detached from hidden_detached).
-    // Recompute rms_norm with grad tracking to get hidden's gradient.
-    if (hidden_normed.grad().defined()) {
-        hidden.set_requires_grad(true);
-        auto hidden_normed_recompute = rms_norm(hidden, final_norm, ctx->rms_eps);
-        hidden_normed_recompute.backward(hidden_normed.grad());
-        // hidden.grad() now has the CE gradient contribution
-    }
+            // Forward — force checkpoint for multi-LoRA (needed for group_inputs)
+            bool use_fused = env_enabled("QWEN36_FUSED_LAYER");
+            TORCH_CHECK(!use_fused,
+                "QWEN36_FUSED_LAYER is disabled until its custom backward preserves the layer graph");
+            ctx->use_checkpoint = true;  // force checkpoint for manual_group_backward
 
-    // emptyCache for CE intermediates when seq>4096.
-    if (hidden_normed.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache();
+            auto t_fwd_start = std::chrono::steady_clock::now();
+            auto hidden = use_fused
+                ? forward_full_fused(ctx, input_ref)
+                : forward_full_checkpoint(ctx, input_ref);
+            auto t_fwd_end = std::chrono::steady_clock::now();
+            double fwd_ms = std::chrono::duration<double, std::milli>(t_fwd_end - t_fwd_start).count();
 
-    return at::tensor({total_loss_val / total_count_val},
-        at::TensorOptions().dtype(at::kFloat).device(hidden.device()));
-}
+            // Batched CE: compute loss with autograd enabled.
+            at::Tensor hidden_grad;
+            at::Tensor loss_hidden;
+            auto t_loss_start = std::chrono::steady_clock::now();
+            {
+                at::AutoGradMode grad_enable(true);
+                // Re-attach hidden to autograd graph
+                hidden = hidden.detach().set_requires_grad(true);
+                // CP owns disjoint sequence slices. Gather them before CE so
+                // every tenant sees the full target row; the gather backward
+                // returns the rank-local gradient slice.
+                loss_hidden = sequence_parallel_loss_gather(ctx, hidden);
+                TORCH_CHECK(!env_enabled("QWEN36_FUSED_CE"),
+                    "QWEN36_FUSED_CE is disabled until its tile gather and gradient normalization are validated");
+                auto loss = compute_loss(
+                    ctx, loss_hidden, input_ref, mask_ref, ctx->vocab_size,
+                    /*independent_samples=*/true,
+                    /*collect_sample_losses=*/loss_numerators.defined());
+                chunk_loss = loss.value.detach();
+                hidden_grad = loss.hidden_grad;
+                if (loss_numerators.defined()) {
+                    TORCH_CHECK(loss.sample_loss_numerators.defined() &&
+                            loss.sample_token_counts.defined() &&
+                            loss.sample_loss_numerators.numel() == n &&
+                            loss.sample_token_counts.numel() == n,
+                        "dynamic per-adapter loss report shape mismatch");
+                    loss_numerators.narrow(0, start, n).copy_(
+                        loss.sample_loss_numerators);
+                    loss_token_counts.narrow(0, start, n).copy_(
+                        loss.sample_token_counts);
+                }
+            }
+            auto t_loss_end = std::chrono::steady_clock::now();
+            double loss_ms = std::chrono::duration<double, std::milli>(t_loss_end - t_loss_start).count();
 
-// ──────────────────────────────────────────────────────────────────────
-// MTP (Multi-Token Prediction) forward + loss
-// ──────────────────────────────────────────────────────────────────────
+            // Backward
+            auto t_bwd_start = std::chrono::steady_clock::now();
+            if (mtp_enabled) {
+                auto mtp_input = hidden.detach().set_requires_grad(true);
+                auto mtp_hidden = mtp_forward(ctx, mtp_input, input_ref);
+                auto mtp_loss = mtp_compute_loss(
+                    ctx, mtp_hidden, input_ref, mask_ref,
+                    /*independent_samples=*/true,
+                    /*collect_sample_losses=*/loss_numerators.defined());
+                mtp_loss.value.backward();
+                TORCH_CHECK(mtp_input.grad().defined(), "MTP did not produce a hidden gradient");
+                auto chunk_gradient_scales = mtp_gradient_scales
+                    .narrow(0, start, n).to(mtp_input.grad().scalar_type())
+                    .reshape({n, 1, 1});
+                hidden_grad.add_(mtp_input.grad() * chunk_gradient_scales);
+                chunk_loss = chunk_loss + mtp_loss.value.detach();
+                if (loss_numerators.defined()) {
+                    TORCH_CHECK(mtp_loss.sample_loss_numerators.defined() &&
+                            mtp_loss.sample_token_counts.defined() &&
+                            mtp_loss.sample_loss_numerators.numel() == n &&
+                            mtp_loss.sample_token_counts.numel() == n,
+                        "dynamic MTP per-adapter loss report shape mismatch");
+                    auto expected_counts = mtp_adapter_token_counts.narrow(
+                        0, start, n);
+                    step_numerics_valid = at::logical_and(
+                        step_numerics_valid,
+                        mtp_loss.sample_token_counts.eq(
+                            expected_counts).all());
+                    loss_numerators.narrow(0, start, n).add_(
+                        mtp_loss.sample_loss_numerators *
+                        mtp_report_scales.narrow(0, start, n));
+                }
+            }
 
-// MTP forward: produce hidden states (not logits) for chunked loss computation.
-// hidden: [batch, seq, hidden] — pre-norm hidden from main model
-// Returns: [batch, seq-1, hidden] — MTP hidden (after final norm, before lm_head)
-static at::Tensor mtp_forward(
-    TrainingContext* ctx,
-    const at::Tensor& hidden,
-    const at::Tensor& input_ids
-) {
-    auto kind = ctx->compute_type;
-    auto embed = *ctx->embed_ptr[0];
+            // independent_samples produces a local per-tenant mean. Restore
+            // each row to a token-sum numerator before A2A backward so the
+            // expert owner can combine unequal source shards correctly.
+            if (ctx->expert_parallel && ctx->nccl_comm &&
+                env_enabled("QWEN36_EP_A2A_SHARDED")) {
+                auto chunk_token_counts = adapter_token_counts
+                    .narrow(0, start, n).reshape({n, 1, 1});
+                hidden_grad.mul_(chunk_token_counts);
+            }
 
-    int64_t seq_len = hidden.size(1);
+            auto chunk_loss_valid = at::logical_and(
+                at::isfinite(chunk_loss), chunk_loss >= 0).all();
+            auto chunk_gradient_valid = at::isfinite(hidden_grad).all();
+            step_numerics_valid = at::logical_and(
+                step_numerics_valid,
+                at::logical_and(chunk_loss_valid, chunk_gradient_valid));
+            total_loss = total_loss + chunk_loss.to(at::kDouble);
+
+            if (!ctx->group_inputs.empty() && !env_enabled("QWEN36_SUBCKPT")) {
+                manual_group_backward(
+                    ctx, sequence_parallel_enabled(ctx) ||
+                        context_parallel_enabled(ctx)
+                        ? sequence_parallel_plain_split(ctx, hidden_grad)
+                        : hidden_grad);
+            } else {
+                loss_hidden.backward(hidden_grad);
+            }
+            // The chunk registry owns moved adapter handles for this chunk.
+            // Harvest now; restoring by canonical index keeps FP32 accumulator
+            // identity intact without copying the full registry maps.
+            harvest_gradient_accumulators(ctx);
+            // No cudaDeviceSynchronize — let GPU pipeline run asynchronously.
+            // The next chunk's CPU prep (LoRA batch, input expand) will overlap
+            // with the tail of this chunk's GPU backward.
+            auto t_bwd_end = std::chrono::steady_clock::now();
+            double bwd_ms = std::chrono::duration<double, std::milli>(t_bwd_end - t_bwd_start).count();
 
-    // hidden[t] + embed[t+1] → predict token t+2 (Megatron convention)
-    auto hidden_shifted = hidden.narrow(1, 0, seq_len - 1);  // [batch, seq-1, hidden]
-    auto embed_next = at::embedding(embed, input_ids.narrow(1, 1, seq_len - 1));  // [batch, seq-1, hidden]
+            if (env_enabled("QWEN36_TRAIN_TRACE")) {
+                fprintf(stderr,
+                    "[train_multi] chunk %ld/%ld: n=%ld loss=device-deferred  "
+                    "fwd=%.0fms loss=%.0fms bwd=%.0fms\n",
+                    (long)(chunk + 1), (long)num_chunks, (long)n,
+                    fwd_ms, loss_ms, bwd_ms);
+            }
+            }
 
-    // RMSNorm both
-    auto h_normed = rms_norm(hidden_shifted, *ctx->mtp_pre_fc_norm_hidden, ctx->rms_eps).to(kind);
-    auto e_normed = rms_norm(embed_next, *ctx->mtp_pre_fc_norm_emb, ctx->rms_eps).to(kind);
+            // Restore the complete registry before the next chunk. Gradients
+            // remain attached to the intrusive tensor handles, so all chunks
+            // can accumulate and the optimizer runs exactly once below.
+            registry_guard.restore();
+
+            if (chunk == num_chunks - 1 && !finalize_only) {
+                bool local_step_valid = false;
+                try {
+                    local_step_valid = step_numerics_valid.item<bool>();
+                } catch (...) {
+                    local_step_valid = false;
+                }
+                const bool global_step_valid =
+                    adapter_collective_all_succeeded(ctx, local_step_valid);
+                TORCH_CHECK(global_step_valid && local_step_valid,
+                    "dynamic LoRA loss, hidden gradient, or MTP token counts "
+                    "became invalid");
+            }
 
-    // Combine: embed first, then hidden → fc projection
-    auto combined = at::cat({e_normed, h_normed}, /*dim=*/-1);
-    auto projected = at::matmul(combined, ctx->mtp_fc->t());  // fc: [hidden, 2*hidden]
+            if (chunk == num_chunks - 1 && !train_only) {
+                // DP gradient synchronization and Adam belong to the logical
+                // multi-tenant step, never to an activation-memory chunk.
+                ++ctx->dynamic_finalizer_count;
+                TORCH_CHECK(!finalize_only || !env_enabled(
+                        "QWEN36_TEST_FAIL_FINALIZER_BEFORE_TOKEN_PREFLIGHT"),
+                    "injected dynamic finalizer failure before token preflight");
+                bool local_token_counts_valid = false;
+                try {
+                    local_token_counts_valid =
+                        at::logical_and(
+                            at::isfinite(adapter_token_counts),
+                            adapter_token_counts >= 0).all().item<bool>();
+                } catch (...) {
+                    local_token_counts_valid = false;
+                }
+                bool global_token_counts_valid = local_token_counts_valid;
+                if ((ctx->nccl_comm && ctx->ep_world_size > 1) ||
+                    (ctx->dp_comm && ctx->dp_world_size > 1) ||
+                    (ctx->tp_comm && ctx->tp_world_size > 1) ||
+                    (ctx->cp_comm && ctx->cp_world_size > 1)) {
+                    global_token_counts_valid =
+                        adapter_collective_all_succeeded(
+                            ctx, local_token_counts_valid);
+                }
+                if (finalize_only && finalizer_phase) *finalizer_phase = 1;
+                TORCH_CHECK(global_token_counts_valid &&
+                        local_token_counts_valid,
+                    "dynamic LoRA token counts must be finite, non-negative, "
+                    "and valid on every distributed rank");
+                std::vector<uint8_t> adapter_has_global_tokens;
+                synchronize_lora_gradients(
+                    ctx, target_mask, 0.0, &adapter_token_counts,
+                    &adapter_has_global_tokens,
+                    /*adapter_token_counts_prevalidated=*/true,
+                    accumulators_are_numerators);
+                TORCH_CHECK(adapter_has_global_tokens.size() ==
+                        ctx->adapters.size(),
+                    "dynamic LoRA global-token activity vector mismatch");
+                const bool inject_dynamic_nan =
+                    env_enabled("QWEN36_TEST_INJECT_DYNAMIC_GRAD_NAN");
+                bool injection_ready = !inject_dynamic_nan;
+                if (inject_dynamic_nan) {
+                    try {
+                        for (size_t adapter_index = 0;
+                             adapter_index < ctx->adapters.size() &&
+                                 !injection_ready;
+                             ++adapter_index) {
+                            if (!adapter_has_global_tokens[adapter_index]) continue;
+                            for (auto& [layer_idx, pairs] :
+                                    ctx->adapters[adapter_index].grad_accum) {
+                                (void)layer_idx;
+                                for (auto& accumulators : pairs) {
+                                    if (accumulators[0].defined() &&
+                                            accumulators[0].numel() > 0) {
+                                        accumulators[0].view({-1})
+                                            .narrow(0, 0, 1)
+                                            .fill_(std::numeric_limits<float>::quiet_NaN());
+                                        injection_ready = true;
+                                        break;
+                                    }
+                                }
+                                if (injection_ready) break;
+                            }
+                        }
+                    } catch (...) {
+                        injection_ready = false;
+                    }
+                }
+                bool local_optimizer_state_ready = true;
+                try {
+                    TORCH_CHECK(!ctx->dynamic_adam_transaction_active,
+                        "dynamic Adam transaction is already active");
+                    ctx->dynamic_adam_transaction_active = true;
+                    shadow_scope.active = true;
+                    for (size_t adapter_index = 0;
+                         adapter_index < ctx->adapters.size();
+                         ++adapter_index) {
+                        if (!adapter_has_global_tokens[adapter_index])
+                            continue;
+                        materialize_dynamic_optimizer_state(
+                            ctx, ctx->adapters[adapter_index]);
+                    }
+                } catch (...) {
+                    local_optimizer_state_ready = false;
+                }
+                bool local_gradients_finite = false;
+                if (injection_ready && local_optimizer_state_ready) {
+                    try {
+                        local_gradients_finite =
+                            dynamic_lora_accumulators_are_finite(
+                                ctx, adapter_has_global_tokens);
+                    } catch (...) {
+                        local_gradients_finite = false;
+                    }
+                }
+                const bool all_gradients_finite =
+                    adapter_collective_all_succeeded(
+                        ctx, local_gradients_finite);
+                TORCH_CHECK(local_optimizer_state_ready &&
+                        local_gradients_finite && all_gradients_finite,
+                    "dynamic LoRA optimizer rejected missing or non-finite "
+                    "accumulated gradients or could not materialize lazy state");
+
+                clip_dynamic_lora_gradients(ctx, adapter_has_global_tokens);
+
+                // Build every Adam result out of place. No live parameter,
+                // optimizer tensor, or tenant clock changes until the unified
+                // launch has passed preflight and CUDA launch validation.
+                at::AutoGradMode guard(false);
+                ctx->lora_cache_valid = false;
+                ctx->lora_batch_valid = false;
+                struct DynamicAdamCommit {
+                    at::Tensor* param;
+                    at::Tensor* accumulator;
+                    at::Tensor* m;
+                    at::Tensor* v;
+                    at::Tensor* next_param;
+                    at::Tensor* next_m;
+                    at::Tensor* next_v;
+                    float lr_scaled;
+                    float eps_scaled;
+                    float beta1;
+                    float beta2;
+                };
+                struct DynamicAdamClockCommit {
+                    TrainingContext::LoRAAdapter* adapter;
+                    int64_t logical_step;
+                };
+                std::vector<DynamicAdamCommit> commits;
+                std::vector<DynamicAdamClockCommit> clock_commits;
+                auto append_commit = [&](at::Tensor& param,
+                                         at::Tensor& accumulator,
+                                         at::Tensor& m,
+                                         at::Tensor& v,
+                                         at::Tensor& next_param,
+                                         at::Tensor& next_m,
+                                         at::Tensor& next_v,
+                                         float lr_scaled,
+                                         float eps_scaled,
+                                         float beta1,
+                                         float beta2) {
+                    if (!param.requires_grad() || !accumulator.defined()) return;
+                    TORCH_CHECK(param.defined() && param.is_cuda() &&
+                            param.is_contiguous() &&
+                            param.scalar_type() == at::kBFloat16,
+                        "dynamic Adam parameter must be contiguous CUDA BF16");
+                    TORCH_CHECK(accumulator.is_cuda() &&
+                            accumulator.is_contiguous() &&
+                            accumulator.scalar_type() == at::kFloat &&
+                            accumulator.sizes() == param.sizes() &&
+                            accumulator.device() == param.device(),
+                        "dynamic Adam accumulator must be matching contiguous CUDA FP32");
+                    TORCH_CHECK(m.defined() && v.defined() && m.is_cuda() &&
+                            v.is_cuda() && m.is_contiguous() && v.is_contiguous() &&
+                            m.scalar_type() == at::kFloat &&
+                            v.scalar_type() == at::kFloat &&
+                            m.sizes() == param.sizes() &&
+                            v.sizes() == param.sizes() &&
+                            m.device() == param.device() &&
+                            v.device() == param.device(),
+                        "dynamic Adam state must be matching contiguous CUDA FP32");
+                    TORCH_CHECK(next_param.defined() && next_param.is_cuda() &&
+                            next_param.is_contiguous() &&
+                            next_param.scalar_type() == at::kBFloat16 &&
+                            next_param.sizes() == param.sizes() &&
+                            next_param.device() == param.device() &&
+                            next_param.requires_grad() == param.requires_grad() &&
+                            next_m.defined() && next_v.defined() &&
+                            next_m.is_cuda() && next_v.is_cuda() &&
+                            next_m.is_contiguous() && next_v.is_contiguous() &&
+                            next_m.scalar_type() == at::kFloat &&
+                            next_v.scalar_type() == at::kFloat &&
+                            next_m.sizes() == param.sizes() &&
+                            next_v.sizes() == param.sizes() &&
+                            next_m.device() == param.device() &&
+                            next_v.device() == param.device(),
+                        "dynamic Adam shadow must match live tensor layout");
+                    TORCH_CHECK(param.numel() > 0 &&
+                            param.numel() <= std::numeric_limits<int>::max(),
+                        "dynamic Adam tensor size is outside fused-kernel range: ",
+                        param.numel());
+                    commits.push_back(DynamicAdamCommit{
+                        &param, &accumulator, &m, &v,
+                        &next_param, &next_m, &next_v,
+                        lr_scaled, eps_scaled, beta1, beta2});
+                };
+                for (size_t adapter_index = 0;
+                     adapter_index < ctx->adapters.size(); ++adapter_index) {
+                    if (!adapter_has_global_tokens[adapter_index]) continue;
+                    auto& adapter = ctx->adapters[adapter_index];
+                    TORCH_CHECK(adapter.optimizer_step >= 0 &&
+                            adapter.optimizer_step <
+                                std::numeric_limits<int64_t>::max(),
+                        "dynamic Adam optimizer clock is outside valid range for adapter ",
+                        adapter.id, ": ", adapter.optimizer_step);
+                    const int64_t logical_step = adapter.optimizer_step + 1;
+                    const double step_f = (double)logical_step;
+                    const double bias_correction1 =
+                        1.0 - std::pow(adapter.optimizer_beta1, step_f);
+                    const double bias_correction2 =
+                        1.0 - std::pow(adapter.optimizer_beta2, step_f);
+                    const double sqrt_bias_correction2 =
+                        std::sqrt(bias_correction2);
+                    const float lr_scaled = (float)(
+                        adapter.optimizer_lr * sqrt_bias_correction2 /
+                        bias_correction1);
+                    const float eps_scaled = (float)(
+                        adapter.optimizer_eps * sqrt_bias_correction2);
+                    TORCH_CHECK(std::isfinite(lr_scaled) &&
+                            std::isfinite(eps_scaled) &&
+                            bias_correction1 > 0.0 &&
+                            bias_correction2 > 0.0,
+                        "dynamic Adam bias correction is invalid for adapter ",
+                        adapter.id, " at logical step ", logical_step);
+                    const size_t adapter_commit_begin = commits.size();
+                    for (auto& [layer_idx, pairs] : adapter.params) {
+                        auto state_it = adapter.adam_state.find(layer_idx);
+                        auto shadow_it = adapter.adam_shadow.find(layer_idx);
+                        auto accum_it = adapter.grad_accum.find(layer_idx);
+                        TORCH_CHECK(state_it != adapter.adam_state.end() &&
+                                shadow_it != adapter.adam_shadow.end() &&
+                                accum_it != adapter.grad_accum.end() &&
+                                state_it->second.size() == pairs.size() &&
+                                shadow_it->second.size() == pairs.size() &&
+                                accum_it->second.size() == pairs.size(),
+                            "dynamic Adam registry layout mismatch for adapter ",
+                            adapter.id, " layer ", layer_idx);
+                        for (size_t i = 0; i < pairs.size(); ++i) {
+                            auto& [a, b] = pairs[i];
+                            auto& [m_a, v_a, m_b, v_b] = state_it->second[i];
+                            auto& [next_a, next_m_a, next_v_a,
+                                   next_b, next_m_b, next_v_b] =
+                                shadow_it->second[i];
+                            auto& [accum_a, accum_b] = accum_it->second[i];
+                            append_commit(
+                                a, accum_a, m_a, v_a,
+                                next_a, next_m_a, next_v_a,
+                                lr_scaled, eps_scaled,
+                                static_cast<float>(adapter.optimizer_beta1),
+                                static_cast<float>(adapter.optimizer_beta2));
+                            append_commit(
+                                b, accum_b, m_b, v_b,
+                                next_b, next_m_b, next_v_b,
+                                lr_scaled, eps_scaled,
+                                static_cast<float>(adapter.optimizer_beta1),
+                                static_cast<float>(adapter.optimizer_beta2));
+                        }
+                    }
+                    if (commits.size() > adapter_commit_begin || allow_pipeline) {
+                        clock_commits.push_back(
+                            DynamicAdamClockCommit{&adapter, logical_step});
+                    }
+                }
+                // All PP stages must rendezvous before any live/shadow swap;
+                // a stage with no local parameters still participates.
+                if (allow_pipeline) {
+                    TORCH_CHECK(adapter_collective_all_succeeded(ctx, true),
+                        "dynamic pipeline Adam commit consensus failed");
+                }
+                if (!commits.empty()) {
+                    TORCH_CHECK(commits.size() <=
+                            static_cast<size_t>(std::numeric_limits<int>::max()),
+                        "dynamic Adam tensor count exceeds fused-kernel range");
+                    std::vector<void*> h_params, h_grads, h_dst_params;
+                    std::vector<float*> h_m, h_v, h_dst_m, h_dst_v;
+                    std::vector<int> h_sizes;
+                    std::vector<float> h_lr_scaled, h_eps_scaled;
+                    std::vector<float> h_beta1, h_beta2;
+                    const size_t tensor_count = commits.size();
+                    h_params.reserve(tensor_count);
+                    h_grads.reserve(tensor_count);
+                    h_m.reserve(tensor_count);
+                    h_v.reserve(tensor_count);
+                    h_dst_params.reserve(tensor_count);
+                    h_dst_m.reserve(tensor_count);
+                    h_dst_v.reserve(tensor_count);
+                    h_sizes.reserve(tensor_count);
+                    h_lr_scaled.reserve(tensor_count);
+                    h_eps_scaled.reserve(tensor_count);
+                    h_beta1.reserve(tensor_count);
+                    h_beta2.reserve(tensor_count);
+                    for (auto& commit : commits) {
+                        h_params.push_back(commit.param->data_ptr());
+                        h_grads.push_back(commit.accumulator->data_ptr());
+                        h_m.push_back((float*)commit.m->data_ptr());
+                        h_v.push_back((float*)commit.v->data_ptr());
+                        h_dst_params.push_back(commit.next_param->data_ptr());
+                        h_dst_m.push_back((float*)commit.next_m->data_ptr());
+                        h_dst_v.push_back((float*)commit.next_v->data_ptr());
+                        h_sizes.push_back((int)commit.param->numel());
+                        h_lr_scaled.push_back(commit.lr_scaled);
+                        h_eps_scaled.push_back(commit.eps_scaled);
+                        h_beta1.push_back(commit.beta1);
+                        h_beta2.push_back(commit.beta2);
+                    }
+                    const int n_params = (int)commits.size();
+                    auto opts_cpu_long = at::TensorOptions().dtype(at::kLong).device(at::kCPU);
+                    auto opts_cpu_int  = at::TensorOptions().dtype(at::kInt).device(at::kCPU);
+                    auto opts_cpu_float = at::TensorOptions().dtype(at::kFloat).device(at::kCPU);
+                    auto params_cpu = at::from_blob(h_params.data(), {n_params}, opts_cpu_long);
+                    auto grads_cpu  = at::from_blob(h_grads.data(),  {n_params}, opts_cpu_long);
+                    auto m_cpu      = at::from_blob(h_m.data(),      {n_params}, opts_cpu_long);
+                    auto v_cpu      = at::from_blob(h_v.data(),      {n_params}, opts_cpu_long);
+                    auto dst_params_cpu = at::from_blob(
+                        h_dst_params.data(), {n_params}, opts_cpu_long);
+                    auto dst_m_cpu = at::from_blob(
+                        h_dst_m.data(), {n_params}, opts_cpu_long);
+                    auto dst_v_cpu = at::from_blob(
+                        h_dst_v.data(), {n_params}, opts_cpu_long);
+                    auto sizes_cpu  = at::from_blob(h_sizes.data(),  {n_params}, opts_cpu_int);
+                    auto lr_cpu = at::from_blob(
+                        h_lr_scaled.data(), {n_params}, opts_cpu_float);
+                    auto eps_cpu = at::from_blob(
+                        h_eps_scaled.data(), {n_params}, opts_cpu_float);
+                    auto beta1_cpu = at::from_blob(
+                        h_beta1.data(), {n_params}, opts_cpu_float);
+                    auto beta2_cpu = at::from_blob(
+                        h_beta2.data(), {n_params}, opts_cpu_float);
+                    ctx->adam_dev_bufs.ensure(n_params, *commits[0].param);
+                    ctx->adam_dev_bufs.params_buf.narrow(0, 0, n_params).copy_(params_cpu);
+                    ctx->adam_dev_bufs.grads_buf.narrow(0, 0, n_params).copy_(grads_cpu);
+                    ctx->adam_dev_bufs.m_buf.narrow(0, 0, n_params).copy_(m_cpu);
+                    ctx->adam_dev_bufs.v_buf.narrow(0, 0, n_params).copy_(v_cpu);
+                    ctx->adam_dev_bufs.dst_params_buf.narrow(0, 0, n_params).copy_(dst_params_cpu);
+                    ctx->adam_dev_bufs.dst_m_buf.narrow(0, 0, n_params).copy_(dst_m_cpu);
+                    ctx->adam_dev_bufs.dst_v_buf.narrow(0, 0, n_params).copy_(dst_v_cpu);
+                    ctx->adam_dev_bufs.sizes_buf.narrow(0, 0, n_params).copy_(sizes_cpu);
+                    ctx->adam_dev_bufs.lr_buf.narrow(0, 0, n_params).copy_(lr_cpu);
+                    ctx->adam_dev_bufs.eps_buf.narrow(0, 0, n_params).copy_(eps_cpu);
+                    ctx->adam_dev_bufs.beta1_buf.narrow(
+                        0, 0, n_params).copy_(beta1_cpu);
+                    ctx->adam_dev_bufs.beta2_buf.narrow(
+                        0, 0, n_params).copy_(beta2_cpu);
+                    auto stream = c10::cuda::getCurrentCUDAStream().stream();
+                    ++ctx->dynamic_adam_launch_count;
+                    launch_fused_adam_multi_out_of_place(
+                        (void**)ctx->adam_dev_bufs.params_buf.data_ptr(),
+                        (void**)ctx->adam_dev_bufs.grads_buf.data_ptr(),
+                        (float**)ctx->adam_dev_bufs.m_buf.data_ptr(),
+                        (float**)ctx->adam_dev_bufs.v_buf.data_ptr(),
+                        (void**)ctx->adam_dev_bufs.dst_params_buf.data_ptr(),
+                        (float**)ctx->adam_dev_bufs.dst_m_buf.data_ptr(),
+                        (float**)ctx->adam_dev_bufs.dst_v_buf.data_ptr(),
+                        (int*)ctx->adam_dev_bufs.sizes_buf.data_ptr(),
+                        (float*)ctx->adam_dev_bufs.lr_buf.data_ptr(),
+                        (float*)ctx->adam_dev_bufs.eps_buf.data_ptr(),
+                        (float*)ctx->adam_dev_bufs.beta1_buf.data_ptr(),
+                        (float*)ctx->adam_dev_bufs.beta2_buf.data_ptr(),
+                        n_params,
+                        (void*)stream);
+                    auto launch_error = cudaGetLastError();
+                    TORCH_CHECK(launch_error == cudaSuccess,
+                        "dynamic transactional Adam launch failed: ",
+                        cudaGetErrorString(launch_error));
+                    const bool inject_failure = env_enabled(
+                        "QWEN36_TEST_FAIL_DYNAMIC_ADAM_BEFORE_COMMIT");
+                    if (inject_failure || env_enabled(
+                            "QWEN36_STRICT_DYNAMIC_ADAM_COMMIT")) {
+                        auto completion_error = cudaStreamSynchronize(stream);
+                        TORCH_CHECK(completion_error == cudaSuccess,
+                            "dynamic transactional Adam execution failed: ",
+                            cudaGetErrorString(completion_error));
+                    }
+                    // Production commits remain asynchronous: the next
+                    // forward consumes the swapped destinations on this same
+                    // stream. Strict mode is available for validation without
+                    // imposing a device barrier on every training step.
+                    TORCH_CHECK(!inject_failure,
+                        "injected dynamic Adam failure before commit");
+                    for (auto& commit : commits) {
+                        std::swap(*commit.param, *commit.next_param);
+                        std::swap(*commit.m, *commit.next_m);
+                        std::swap(*commit.v, *commit.next_v);
+                    }
+                }
+                for (const auto& clock_commit : clock_commits) {
+                    clock_commit.adapter->optimizer_step =
+                        clock_commit.logical_step;
+                }
+            }
 
-    // MTP layers (full attention + MoE/dense, no LoRA)
-    at::Tensor h = projected;
-    int64_t num_mtp_layers = (int64_t)ctx->mtp_layer_configs.size();
-    for (int64_t i = 0; i < num_mtp_layers; i++) {
-        int64_t w_offset = 0;
-        for (int64_t j = 0; j < i; j++)
-            w_offset += weight_count_for_layer(ctx->mtp_layer_configs[j]);
-        int64_t w_count = weight_count_for_layer(ctx->mtp_layer_configs[i]);
-        std::vector<at::Tensor*> layer_w(ctx->mtp_layer_weights.begin() + w_offset,
-                                         ctx->mtp_layer_weights.begin() + w_offset + w_count);
-        // MTP processes seq-1 tokens — slice attention mask's last dim to match
-        auto mtp_mask = ctx->attention_mask.defined()
-            ? ctx->attention_mask.narrow(-1, 0, h.size(1))
-            : at::Tensor();
-        h = forward_single_layer(ctx, h, layer_w.data(), &ctx->mtp_layer_configs[i],
-            ctx->num_layers + i, kind, mtp_mask);
-    }
+        }
 
-    // Final norm only — return hidden, not logits
-    return rms_norm(h, *ctx->mtp_norm, ctx->rms_eps).to(kind);
+        if (!train_only) clear_gradient_accumulators(ctx);
+        if (!finalize_only && loss_numerators_out && token_counts_out) {
+            *loss_numerators_out = loss_numerators;
+            *token_counts_out = loss_token_counts;
+        }
+        accumulation_guard.disarmed = true;
+        if (finalize_only || loss_numerators.defined()) return 0.0;
+        return total_loss.item<double>() / total_adapters;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[train_multi] FAILED: %s\n", e.what());
+        return -1.0;
+    } catch (...) {
+        fprintf(stderr, "[train_multi] FAILED: unknown exception\n");
+        return -1.0;
+    }
 }
 
-// MTP loss: chunked matmul + cross-entropy, weighted by 0.5
-// mtp_hidden[t] (from hidden[t] + embed[t+1]) predicts token t+2 (Megatron convention)
-// No full logits tensor — chunked matmul + fused CE
-static at::Tensor mtp_compute_loss(
-    TrainingContext* ctx,
-    const at::Tensor& mtp_hidden,
-    const at::Tensor& input_ids,
-    const at::Tensor& target_mask
+__attribute__((visibility("default"))) double qwen36_train_multi_lora(
+    void* ctx_ptr,
+    void* input_ids_ptr,
+    void* target_mask_ptr,
+    void* attention_mask_ptr,
+    int32_t n_total,
+    int32_t lora_rank
 ) {
-    int64_t vocab_size = ctx->vocab_size;
-    int64_t seq_len = input_ids.size(1);
-    auto lm_head = *ctx->lm_head_ptr[0];
-
-    // MTP hidden: [batch, seq-1, hidden], drop last → predict t+2
-    int64_t n_tokens = seq_len - 2;
-    auto hidden_flat = mtp_hidden.narrow(1, 0, n_tokens).reshape({-1, mtp_hidden.size(2)});
-    auto shifted_targets = input_ids.narrow(1, 2, n_tokens).reshape({-1});
-    auto shifted_mask = target_mask.narrow(1, 2, n_tokens).reshape({-1});
-
-    // Chunked matmul + cross-entropy
-    int64_t total_tokens = shifted_targets.size(0);
-    int64_t chunk_size = 4096;  // smaller chunks = less peak memory
-    int64_t num_chunks = (total_tokens + chunk_size - 1) / chunk_size;
-
-    auto total_loss = at::zeros({1}, at::TensorOptions().dtype(at::kFloat).device(mtp_hidden.device()));
-    auto total_count = shifted_mask.sum().clamp_min(1.0);
+    return qwen36_train_multi_lora_impl(
+        ctx_ptr, input_ids_ptr, target_mask_ptr, attention_mask_ptr,
+        n_total, lora_rank, DynamicMultiLoraMode::TrainAndFinalize);
+}
 
-    for (int64_t c = 0; c < num_chunks; c++) {
-        int64_t start = c * chunk_size;
-        int64_t end = std::min(start + chunk_size, total_tokens);
-        int64_t n = end - start;
+// Train only the requested dynamic tenants. The existing train_multi_lora
+// implementation already batches activation-level LoRA projections and owns
+// the logical Adam boundary; this wrapper scopes its adapter registry to the
+// selected IDs and restores the original order on every exit.
+__attribute__((visibility("default"))) double qwen36_train_multi_lora_selected(
+    void* ctx_ptr,
+    void* input_ids_ptr,
+    void* target_mask_ptr,
+    void* attention_mask_ptr,
+    const int64_t* adapter_ids,
+    int32_t n_adapters,
+    int32_t lora_rank
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    std::vector<TrainingContext::LoRAAdapter> original;
+    std::vector<TrainingContext::LoRAAdapter> selected;
+    std::vector<TrainingContext::LoRAAdapter> merged;
+    std::vector<size_t> selected_indexes;
+    std::vector<uint8_t> moved;
+    bool registry_detached = false;
+    bool selected_installed = false;
+    bool adapter_id_index_was_valid = false;
+    auto restore_registry = [&]() {
+        if (!ctx || !registry_detached) return;
+        if (selected_installed) {
+            selected.swap(ctx->adapters);
+            selected_installed = false;
+        }
+        for (size_t i = 0; i < original.size(); ++i) {
+            if (!moved[i]) {
+                merged.push_back(std::move(original[i]));
+                continue;
+            }
+            auto selected_it = std::find(
+                selected_indexes.begin(), selected_indexes.end(), i);
+            TORCH_CHECK(selected_it != selected_indexes.end(),
+                "selected adapter index disappeared: ", i);
+            const auto selected_index = static_cast<size_t>(
+                selected_it - selected_indexes.begin());
+            merged.push_back(std::move(selected[selected_index]));
+        }
+        ctx->adapters.swap(merged);
+        ctx->adapter_id_index_valid = adapter_id_index_was_valid;
+        registry_detached = false;
+    };
+    try {
+        TORCH_CHECK(ctx,
+            "selected multi-LoRA requires a valid training context");
+        validate_dynamic_context_health(ctx);
+        validate_native_execution_topology_collective(ctx);
+        validate_adapter_collective_registry(
+            ctx, adapter_ids, n_adapters, lora_rank, false);
+        TORCH_CHECK(adapter_ids && n_adapters > 0,
+            "selected multi-LoRA requires at least one adapter ID");
+        const auto original_count = ctx->adapters.size();
+        original.reserve(original_count);
+        selected.reserve(n_adapters);
+        merged.reserve(original_count);
+        selected_indexes.reserve(n_adapters);
+        moved.resize(original_count, 0);
+        require_canonical_adapter_index(ctx);
+        for (int32_t i = 0; i < n_adapters; ++i) {
+            TORCH_CHECK(adapter_ids[i] > 0,
+                "selected adapter IDs must be positive");
+            size_t index = 0;
+            TORCH_CHECK(find_canonical_adapter_index(
+                    ctx, adapter_ids[i], index),
+                "unknown selected adapter ID: ", adapter_ids[i]);
+            TORCH_CHECK(!moved[index],
+                "duplicate selected adapter ID: ", adapter_ids[i]);
+            moved[index] = 1;
+            selected_indexes.push_back(index);
+        }
+        adapter_id_index_was_valid = ctx->adapter_id_index_valid;
+        ctx->adapter_id_index_valid = false;
+        original.swap(ctx->adapters);
+        registry_detached = true;
+        for (int32_t i = 0; i < n_adapters; ++i) {
+            selected.push_back(
+                std::move(original[selected_indexes[static_cast<size_t>(i)]]));
+        }
+        ctx->adapters.swap(selected);
+        selected_installed = true;
+        const double loss = qwen36_train_multi_lora(
+            ctx_ptr, input_ids_ptr, target_mask_ptr, attention_mask_ptr,
+            n_adapters, lora_rank);
+        restore_registry();
+        if (std::isfinite(loss) && loss >= 0.0)
+            page_cold_dynamic_adam_state(
+                ctx, adapter_ids, n_adapters);
+        return loss;
+    } catch (const std::exception& e) {
+        try {
+            restore_registry();
+        } catch (...) {
+            if (ctx) ctx->poisoned = true;
+            fprintf(stderr,
+                "[train_multi_selected] registry restore FAILED; context poisoned\n");
+        }
+        fprintf(stderr, "[train_multi_selected] FAILED: %s\n", e.what());
+        return -1.0;
+    } catch (...) {
+        try {
+            restore_registry();
+        } catch (...) {
+            if (ctx) ctx->poisoned = true;
+            fprintf(stderr,
+                "[train_multi_selected] registry restore FAILED; context poisoned\n");
+        }
+        fprintf(stderr, "[train_multi_selected] FAILED: unknown exception\n");
+        return -1.0;
+    }
+}
 
-        auto chunk_hidden = hidden_flat.narrow(0, start, n);
-        auto chunk_logits = at::matmul(chunk_hidden, lm_head.t());  // [n, vocab]
-        auto chunk_targets = shifted_targets.narrow(0, start, n);
-        auto chunk_mask = shifted_mask.narrow(0, start, n);
+static std::string dynamic_adapter_group_key(
+    const TrainingContext::LoRAAdapter& adapter
+) {
+    std::ostringstream key;
+    key << "rank=" << adapter.rank;
+    key << "|all_layers=" << adapter.all_target_layers;
+    key << "|global_layers=";
+    for (const auto layer : adapter.global_target_layers) key << layer << ",";
+    key << "|modules=";
+    for (const auto& module : adapter.target_modules) key << module << ",";
+    return key.str();
+}
 
-        auto per_token_loss = at::cross_entropy_loss(
-            chunk_logits.to(at::kFloat), chunk_targets,
-            /*weight=*/at::Tensor(), /*reduction=*/at::Reduction::None,
-            /*ignore_index=*/-100, /*label_smoothing=*/0.0
-        );
-        auto masked_loss = per_token_loss * chunk_mask.to(at::kFloat);
-        total_loss += masked_loss.sum();
+static void rollback_dynamic_adapter_commit(
+    TrainingContext::LoRAAdapter& adapter,
+    int64_t optimizer_step
+) {
+    if (adapter.optimizer_step == optimizer_step) return;
+    TORCH_CHECK(adapter.optimizer_step == optimizer_step + 1,
+        "heterogeneous rollback observed unexpected optimizer clock for adapter ",
+        adapter.id, ": expected ", optimizer_step + 1,
+        " got ", adapter.optimizer_step);
+    for (auto& [layer_idx, pairs] : adapter.params) {
+        auto state_it = adapter.adam_state.find(layer_idx);
+        auto shadow_it = adapter.adam_shadow.find(layer_idx);
+        TORCH_CHECK(state_it != adapter.adam_state.end() &&
+                shadow_it != adapter.adam_shadow.end() &&
+                state_it->second.size() == pairs.size() &&
+                shadow_it->second.size() == pairs.size(),
+            "heterogeneous rollback registry mismatch for adapter ",
+            adapter.id, " layer ", layer_idx);
+        for (size_t pair_idx = 0; pair_idx < pairs.size(); ++pair_idx) {
+            auto& [a, b] = pairs[pair_idx];
+            auto& [m_a, v_a, m_b, v_b] = state_it->second[pair_idx];
+            auto& [old_a, old_m_a, old_v_a,
+                   old_b, old_m_b, old_v_b] = shadow_it->second[pair_idx];
+            if (a.requires_grad()) {
+                TORCH_CHECK(m_a.defined() && v_a.defined() &&
+                        old_a.defined() && old_m_a.defined() &&
+                        old_v_a.defined(),
+                    "dynamic rollback observed unmaterialized A state for adapter ",
+                    adapter.id, " layer ", layer_idx,
+                    " pair ", pair_idx);
+                std::swap(a, old_a);
+                std::swap(m_a, old_m_a);
+                std::swap(v_a, old_v_a);
+            }
+            if (b.requires_grad()) {
+                TORCH_CHECK(m_b.defined() && v_b.defined() &&
+                        old_b.defined() && old_m_b.defined() &&
+                        old_v_b.defined(),
+                    "dynamic rollback observed unmaterialized B state for adapter ",
+                    adapter.id, " layer ", layer_idx,
+                    " pair ", pair_idx);
+                std::swap(b, old_b);
+                std::swap(m_b, old_m_b);
+                std::swap(v_b, old_v_b);
+            }
+        }
     }
-
-    return (total_loss / total_count) * 0.5;
+    adapter.optimizer_step = optimizer_step;
 }
 
-// ──────────────────────────────────────────────────────────────────────
-// C FFI
-// ──────────────────────────────────────────────────────────────────────
+static at::Tensor global_dynamic_adapter_losses(
+    TrainingContext* ctx,
+    at::Tensor loss_numerators,
+    at::Tensor token_counts
+) {
+    TORCH_CHECK(loss_numerators.defined() && token_counts.defined() &&
+            loss_numerators.dim() == 1 &&
+            loss_numerators.sizes() == token_counts.sizes(),
+        "dynamic loss report tensors must be matching vectors");
+    bool local_valid = false;
+    try {
+        local_valid = at::logical_and(
+                at::isfinite(loss_numerators),
+                at::logical_and(at::isfinite(token_counts), token_counts >= 0))
+            .all().item<bool>();
+    } catch (...) {
+        local_valid = false;
+    }
+    TORCH_CHECK(adapter_collective_all_succeeded(ctx, local_valid) && local_valid,
+        "dynamic loss report contains invalid numerator or token count");
+
+    auto packed = at::cat({loss_numerators, token_counts}, 0).contiguous();
+    auto reduce_source_axis = [&](ncclComm_t communicator, const char* axis) {
+        auto reduced = at::empty_like(packed);
+        auto stream = c10::cuda::getCurrentCUDAStream(
+            packed.device().index()).stream();
+        auto error = ncclAllReduce(
+            packed.data_ptr<float>(), reduced.data_ptr<float>(), packed.numel(),
+            ncclFloat, ncclSum, communicator, stream);
+        TORCH_CHECK(error == ncclSuccess,
+            "dynamic loss ", axis, " reduction failed: ",
+            ncclGetErrorString(error));
+        packed = reduced;
+    };
+    const bool sharded_ep = ctx->expert_parallel && ctx->nccl_comm &&
+        env_enabled("QWEN36_EP_A2A_SHARDED");
+    if (sharded_ep) {
+        reduce_source_axis(
+            reinterpret_cast<ncclComm_t>(ctx->nccl_comm), "sharded EP");
+    }
+    if (ctx->data_parallel && ctx->dp_comm && ctx->dp_world_size > 1) {
+        reduce_source_axis(
+            reinterpret_cast<ncclComm_t>(ctx->dp_comm), "DP");
+    }
 
-extern "C" {
+    const int64_t count = loss_numerators.numel();
+    auto global_numerators = packed.narrow(0, 0, count);
+    auto global_counts = packed.narrow(0, count, count);
+    return at::where(
+        global_counts > 0,
+        global_numerators / global_counts.clamp_min(1.0),
+        at::zeros_like(global_numerators));
+}
 
-// Create training context — called once at startup
-// lora_rank: LoRA rank (from config)
-// target_layers: array of layer indices to apply LoRA (nullptr = all layers)
-// num_target_layers: length of target_layers array
-__attribute__((visibility("default"))) void* qwen36_create_training_context(
-    void** weight_ptrs, int64_t num_weight_ptrs,
-    void* embed_ptr, void* final_norm_ptr, void* lm_head_ptr,
-    void* layer_configs_ptr, int64_t num_layers,
-    int32_t compute_type,
-    double lora_scaling, double lr, double beta1, double beta2, double eps,
-    int64_t vocab_size, double rms_eps,
-    int64_t lora_rank,
-    const int64_t* target_layers, int64_t num_target_layers
+// Finish a selected dynamic PP window while returning one normalized loss per
+// selected tenant. Loss statistics are produced only on the last PP stage;
+// reduce them across PP before the existing DP/EP reductions in the helper.
+extern "C" __attribute__((visibility("default"))) int32_t
+qwen36_pipeline_finish_dynamic_report_v1(
+    void* ctx_ptr,
+    int32_t apply_optimizer,
+    Qwen36PipelineResultV1* result,
+    double* adapter_losses,
+    int32_t adapter_loss_capacity
 ) {
     try {
-        auto* ctx = new TrainingContext();
-        ctx->compute_type = static_cast<at::ScalarType>(compute_type);
-        ctx->lr = lr; ctx->beta1 = beta1; ctx->beta2 = beta2; ctx->eps = eps;
-        ctx->vocab_size = vocab_size; ctx->rms_eps = rms_eps;
-        ctx->step_count = 0; ctx->lora_scaling = lora_scaling;
-        ctx->num_layers = num_layers;
-        ctx->use_checkpoint = false; ctx->group_size = 4;
-
-        // Store weight pointers
-        auto** wp = reinterpret_cast<at::Tensor**>(weight_ptrs);
-        for (int64_t i = 0; i < num_weight_ptrs; i++) {
-            ctx->weight_ptrs.push_back(wp[i]);
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx && ctx->pp_window.active &&
+                ctx->pp_window.dynamic_lora,
+            "dynamic pipeline report requires an active dynamic window");
+        auto& window = ctx->pp_window;
+        const auto selected_count =
+            static_cast<int32_t>(window.selected_adapter_indices.size());
+        TORCH_CHECK(adapter_losses && adapter_loss_capacity >= selected_count,
+            "dynamic pipeline report output buffer is too small");
+        TORCH_CHECK(window.loss_numerators.defined() &&
+                window.loss_token_counts.defined() &&
+                window.loss_numerators.numel() == selected_count &&
+                window.loss_token_counts.numel() == selected_count,
+            "dynamic pipeline report statistics are unavailable");
+        bool local_stats_valid = false;
+        try {
+            local_stats_valid = at::logical_and(
+                    at::isfinite(window.loss_numerators),
+                    at::logical_and(at::isfinite(window.loss_token_counts),
+                        window.loss_token_counts >= 0))
+                .all().item<bool>();
+        } catch (...) {
+            local_stats_valid = false;
         }
-        ctx->embed_ptr.push_back(reinterpret_cast<at::Tensor*>(embed_ptr));
-        ctx->final_norm_ptr.push_back(reinterpret_cast<at::Tensor*>(final_norm_ptr));
-        ctx->lm_head_ptr.push_back(reinterpret_cast<at::Tensor*>(lm_head_ptr));
-
-        // Copy layer configs
-        auto* lcfgs = reinterpret_cast<LayerConfig*>(layer_configs_ptr);
-        for (int64_t i = 0; i < num_layers; i++) {
-            ctx->layer_configs.push_back(lcfgs[i]);
+        TORCH_CHECK(pipeline_control_all_succeeded(ctx, local_stats_valid) &&
+                local_stats_valid,
+            "dynamic pipeline loss statistics are invalid on one or more PP stages");
+
+        auto global_numerators = window.loss_numerators;
+        auto global_counts = window.loss_token_counts;
+        auto comm = ctx->pp_control_comm ? ctx->pp_control_comm : ctx->pp_comm;
+        auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream();
+        auto pp_numerators = at::empty_like(global_numerators);
+        auto pp_counts = at::empty_like(global_counts);
+        TORCH_CHECK(ncclAllReduce(
+                global_numerators.data_ptr<float>(),
+                pp_numerators.data_ptr<float>(), global_numerators.numel(),
+                ncclFloat, ncclSum, comm, stream) == ncclSuccess &&
+            ncclAllReduce(
+                global_counts.data_ptr<float>(), pp_counts.data_ptr<float>(),
+                global_counts.numel(), ncclFloat, ncclSum, comm, stream) ==
+                ncclSuccess,
+            "dynamic pipeline loss statistics PP reduction failed");
+        auto losses = global_dynamic_adapter_losses(
+            ctx, pp_numerators, pp_counts).to(
+                at::TensorOptions().device(at::kCPU).dtype(at::kDouble));
+        std::memcpy(adapter_losses, losses.data_ptr<double>(),
+            static_cast<size_t>(selected_count) * sizeof(double));
+        return qwen36_pipeline_finish_v1(ctx_ptr, apply_optimizer, result);
+    } catch (const std::exception& error) {
+        fprintf(stderr, "[pipeline_window] dynamic report FAILED: %s\n",
+            error.what());
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        if (ctx) {
+            clear_gradient_accumulators(ctx);
+            pipeline_window_reset(ctx);
         }
-
-        // Build target layer set
-        std::set<int64_t> target_set;
-        if (target_layers && num_target_layers > 0) {
-            for (int64_t j = 0; j < num_target_layers; j++)
-                target_set.insert(target_layers[j]);
+        return -1;
+    } catch (...) {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        if (ctx) {
+            clear_gradient_accumulators(ctx);
+            pipeline_window_reset(ctx);
         }
+        return -1;
+    }
+}
 
-        // Create LoRA parameters for target layers only
-        int64_t offset = 0;
-        auto kind = ctx->compute_type;
-        for (int64_t i = 0; i < num_layers; i++) {
-            int64_t lora_count = (ctx->layer_configs[i].layer_type == 0) ? 4 : 3;
-            ctx->lora_layer_offset.push_back(offset);
-
-            if (target_set.find(i) == target_set.end()) {
-                // Not a target layer — no LoRA params, offset stays same
-                continue;
+// Heterogeneous selected training pads each active LoRA rank to the largest
+// rank needed by that projection and inserts zero A/B slots for adapters that
+// do not target it. This keeps the activation batch aligned with request order
+// so all selected tenants share one forward/backward and one transactional
+// synchronization/Adam finalizer.
+static double qwen36_train_multi_lora_selected_impl(
+    void* ctx_ptr,
+    void* input_ids_ptr,
+    void* target_mask_ptr,
+    void* attention_mask_ptr,
+    const int64_t* adapter_ids,
+    int32_t n_adapters,
+    double* aggregate_loss_out,
+    double* adapter_losses_out,
+    int32_t adapter_loss_capacity
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    std::vector<TrainingContext::LoRAAdapter> original;
+    std::vector<TrainingContext::LoRAAdapter> selected;
+    std::vector<size_t> selected_indexes;
+    std::vector<int64_t> original_steps;
+    bool recovery_failed = false;
+    bool registry_detached = false;
+    bool selected_registry_installed = false;
+    bool adapter_id_index_was_valid = false;
+    at::Tensor selected_loss_numerators;
+    at::Tensor selected_token_counts;
+    at::Tensor global_adapter_loss_values;
+    auto restore_registry = [&]() {
+        if (!ctx || !registry_detached) return;
+        TORCH_CHECK(selected.size() == selected_indexes.size(),
+            "heterogeneous selected adapter restore count mismatch");
+        for (size_t selected_index = 0;
+             selected_index < selected.size(); ++selected_index) {
+            const size_t original_index = selected_indexes[selected_index];
+            TORCH_CHECK(original_index < original.size(),
+                "heterogeneous selected adapter restore index is out of range: ",
+                original_index);
+            original[original_index] = std::move(selected[selected_index]);
+        }
+        ctx->adapters.clear();
+        ctx->adapters.swap(original);
+        ctx->adapter_id_index_valid = adapter_id_index_was_valid;
+        registry_detached = false;
+    };
+    auto rollback = [&]() {
+        const size_t rollback_count = std::min(
+            selected.size(), original_steps.size());
+        for (size_t index = 0; index < rollback_count; ++index) {
+            rollback_dynamic_adapter_commit(
+                selected[index], original_steps[index]);
+        }
+    };
+    auto recover = [&]() {
+        try {
+            if (ctx && registry_detached && selected_registry_installed) {
+                selected.swap(ctx->adapters);
+                selected_registry_installed = false;
+            } else if (ctx && registry_detached && !ctx->adapters.empty()) {
+                // A failing group returns with its scoped registry installed.
+                for (auto& adapter : ctx->adapters) {
+                    auto selected_it = std::find_if(
+                        selected.begin(), selected.end(),
+                        [&](const auto& item) { return item.id == adapter.id; });
+                    if (selected_it != selected.end())
+                        *selected_it = std::move(adapter);
+                }
+                ctx->adapters.clear();
+            }
+        } catch (const std::exception& e) {
+            recovery_failed = true;
+            fprintf(stderr,
+                "[train_multi_selected_v2] group recovery FAILED: %s\n",
+                e.what());
+        } catch (...) {
+            recovery_failed = true;
+            fprintf(stderr,
+                "[train_multi_selected_v2] group recovery FAILED\n");
+        }
+        try {
+            if (ctx && registry_detached) {
+                for (auto& adapter : selected)
+                    clear_adapter_gradient_accumulators(adapter);
+                ctx->group_inputs.clear();
+                ctx->group_outputs.clear();
+                ctx->lora_cache_valid = false;
+                ctx->lora_batch_valid = false;
+                ctx->accumulation_active = false;
+                ctx->accumulated_token_weight = 0.0;
+            }
+        } catch (const std::exception& e) {
+            recovery_failed = true;
+            fprintf(stderr,
+                "[train_multi_selected_v2] gradient cleanup FAILED: %s\n",
+                e.what());
+        } catch (...) {
+            recovery_failed = true;
+            fprintf(stderr,
+                "[train_multi_selected_v2] gradient cleanup FAILED\n");
+        }
+        try {
+            rollback();
+        } catch (const std::exception& e) {
+            recovery_failed = true;
+            fprintf(stderr,
+                "[train_multi_selected_v2] rollback FAILED: %s\n", e.what());
+        } catch (...) {
+            recovery_failed = true;
+            fprintf(stderr, "[train_multi_selected_v2] rollback FAILED\n");
+        }
+        release_dynamic_adam_shadows(ctx, selected);
+        if (ctx) ctx->dynamic_adam_transaction_active = false;
+        try {
+            restore_registry();
+        } catch (const std::exception& e) {
+            recovery_failed = true;
+            fprintf(stderr,
+                "[train_multi_selected_v2] registry restore FAILED: %s\n",
+                e.what());
+        } catch (...) {
+            recovery_failed = true;
+            fprintf(stderr,
+                "[train_multi_selected_v2] registry restore FAILED\n");
+        }
+        if (env_enabled("QWEN36_TEST_POISON_DYNAMIC_RECOVERY"))
+            recovery_failed = true;
+        if (ctx && recovery_failed) {
+            ctx->poisoned = true;
+            fprintf(stderr,
+                "[train_multi_selected_v2] context poisoned after failed recovery\n");
+        }
+    };
+    try {
+        TORCH_CHECK(ctx,
+            "heterogeneous selected training requires a context");
+        validate_native_execution_topology_collective(ctx);
+        const bool local_healthy = !ctx->poisoned;
+        const bool local_request_valid = adapter_ids && n_adapters > 0;
+        const bool report_requested = aggregate_loss_out || adapter_losses_out;
+        const bool local_report_valid =
+            (!report_requested && !aggregate_loss_out && !adapter_losses_out) ||
+            (aggregate_loss_out && adapter_losses_out &&
+                adapter_loss_capacity >= n_adapters);
+        const bool local_accumulation_clear = !ctx->accumulation_active &&
+            ctx->accumulated_token_weight == 0.0;
+        const auto preflight = adapter_collective_min_flags(ctx, {
+            local_healthy ? 1 : 0,
+            local_request_valid ? 1 : 0,
+            report_requested ? 1 : 0,
+            report_requested ? 0 : 1,
+            local_report_valid ? 1 : 0,
+            local_accumulation_clear ? 1 : 0,
+        });
+        TORCH_INTERNAL_ASSERT(preflight.size() == 6);
+        if (!preflight[0]) ctx->poisoned = true;
+        TORCH_CHECK(local_healthy && preflight[0],
+            "dynamic multi-LoRA context is poisoned; restart the worker");
+        TORCH_CHECK(local_request_valid && preflight[1],
+            "heterogeneous selected training requires rank-consistent "
+            "adapter IDs and a positive adapter count");
+        TORCH_CHECK(preflight[2] || preflight[3],
+            "heterogeneous selected loss report capability differs across "
+            "distributed ranks");
+        TORCH_CHECK(local_report_valid && preflight[4],
+            "heterogeneous selected loss report requires aggregate and adapter "
+            "outputs with capacity for every selected adapter");
+        TORCH_CHECK(local_accumulation_clear && preflight[5],
+            "cannot start heterogeneous multi-LoRA while a fixed or dynamic "
+            "gradient accumulation window is pending");
+        validate_adapter_collective_registry(
+            ctx, adapter_ids, n_adapters, 0, false);
+        auto* input_ids_tensor = reinterpret_cast<at::Tensor*>(input_ids_ptr);
+        auto* target_mask_tensor = reinterpret_cast<at::Tensor*>(target_mask_ptr);
+        bool local_input_valid = input_ids_tensor && target_mask_tensor;
+        if (local_input_valid) {
+            local_input_valid = input_ids_tensor->dim() == 2 &&
+                target_mask_tensor->dim() == 2 &&
+                input_ids_tensor->sizes() == target_mask_tensor->sizes();
+        }
+        int64_t input_batch = 0;
+        if (local_input_valid) {
+            input_batch = input_ids_tensor->size(0);
+            local_input_valid = input_batch == 1 || input_batch == n_adapters;
+        }
+        at::Tensor attention_mask;
+        if (local_input_valid && attention_mask_ptr) {
+            attention_mask = *reinterpret_cast<at::Tensor*>(attention_mask_ptr);
+            local_input_valid =
+                attention_mask.sizes() == input_ids_tensor->sizes();
+        }
+        const bool global_input_valid = adapter_collective_all_succeeded(
+            ctx, local_input_valid);
+        TORCH_CHECK(global_input_valid && local_input_valid,
+            "heterogeneous multi-LoRA inputs must be rank-consistent [batch, seq] "
+            "tensors with batch 1 or adapter count");
+        auto& input_ids = *input_ids_tensor;
+        auto& target_mask = *target_mask_tensor;
+
+        const size_t original_count = ctx->adapters.size();
+        require_canonical_adapter_index(ctx);
+        std::vector<size_t> requested_indexes;
+        std::vector<uint8_t> requested(original_count, 0);
+        requested_indexes.reserve(n_adapters);
+        for (int32_t request_index = 0;
+             request_index < n_adapters; ++request_index) {
+            TORCH_CHECK(adapter_ids[request_index] > 0,
+                "heterogeneous selected adapter IDs must be positive");
+            size_t index = 0;
+            TORCH_CHECK(find_canonical_adapter_index(
+                    ctx, adapter_ids[request_index], index),
+                "unknown heterogeneous selected adapter ID: ",
+                adapter_ids[request_index]);
+            TORCH_CHECK(!requested[index],
+                "duplicate heterogeneous selected adapter ID: ",
+                adapter_ids[request_index]);
+            requested[index] = 1;
+            requested_indexes.push_back(index);
+        }
+        original.reserve(original_count);
+        selected.reserve(n_adapters);
+        selected_indexes.reserve(n_adapters);
+        original_steps.reserve(requested_indexes.size());
+        adapter_id_index_was_valid = ctx->adapter_id_index_valid;
+        ctx->adapter_id_index_valid = false;
+        original.swap(ctx->adapters);
+        registry_detached = true;
+        for (const size_t index : requested_indexes) {
+            selected_indexes.push_back(index);
+            original_steps.push_back(original[index].optimizer_step);
+            selected.push_back(std::move(original[index]));
+        }
+        double train_loss = 0.0;
+        if (env_enabled("QWEN36_HETERO_PADDED_BATCH", true)) {
+            int64_t maximum_rank = 0;
+            for (const auto& adapter : selected)
+                maximum_rank = std::max(maximum_rank, adapter.rank);
+            TORCH_CHECK(maximum_rank > 0 &&
+                    maximum_rank <= std::numeric_limits<int32_t>::max(),
+                "heterogeneous adapter rank is outside the native trainer range: ",
+                maximum_rank);
+
+            ctx->adapters.swap(selected);
+            selected_registry_installed = true;
+            const bool saved_padding_mode = ctx->pad_heterogeneous_lora_batch;
+            ctx->pad_heterogeneous_lora_batch = true;
+            train_loss = qwen36_train_multi_lora_impl(
+                ctx, input_ids_ptr, target_mask_ptr, attention_mask_ptr,
+                n_adapters, static_cast<int32_t>(maximum_rank),
+                DynamicMultiLoraMode::TrainOnly, nullptr,
+                report_requested ? &selected_loss_numerators : nullptr,
+                report_requested ? &selected_token_counts : nullptr);
+            ctx->pad_heterogeneous_lora_batch = saved_padding_mode;
+            selected.swap(ctx->adapters);
+            selected_registry_installed = false;
+
+            bool local_batch_succeeded =
+                std::isfinite(train_loss) && train_loss >= 0.0;
+            const char* fail_after = std::getenv(
+                "QWEN36_TEST_FAIL_HETERO_GROUP_AFTER");
+            if (fail_after && fail_after[0] != '\0') {
+                char* end = nullptr;
+                const long requested = std::strtol(fail_after, &end, 10);
+                local_batch_succeeded = local_batch_succeeded &&
+                    end && *end == '\0' && requested > 0 && requested != 1;
+            }
+            TORCH_CHECK(adapter_collective_all_succeeded(
+                    ctx, local_batch_succeeded),
+                "heterogeneous adapter batch failed on at least one "
+                "distributed rank");
+        } else {
+            if (report_requested) {
+                selected_loss_numerators = at::zeros({n_adapters},
+                    at::TensorOptions().dtype(at::kFloat).device(input_ids.device()));
+                selected_token_counts = at::zeros_like(selected_loss_numerators);
+            }
+            std::map<std::string, std::vector<size_t>> groups;
+            for (size_t index = 0; index < selected.size(); ++index) {
+                groups[dynamic_adapter_group_key(selected[index])]
+                    .push_back(index);
             }
 
-            // Get base weight shapes from the weight pointers
-            int64_t w_offset = 0;
-            for (int64_t j = 0; j < i; j++)
-                w_offset += weight_count_for_layer(ctx->layer_configs[j]);
-
-            if (ctx->layer_configs[i].layer_type == 0) {
-                // Full attention: q_proj, k_proj, v_proj, o_proj
-                int64_t proj_indices[] = {2, 4, 6, 7};  // q, k, v, o
-                for (int k = 0; k < 4; k++) {
-                    auto* base = ctx->weight_ptrs[w_offset + proj_indices[k]];
-                    int64_t out_f = base->size(0), in_f = base->size(1);
-                    auto a = at::randn({lora_rank, in_f}, at::TensorOptions().dtype(ctx->compute_type).device(base->device())) * 0.01;
-                    auto b = at::zeros({out_f, lora_rank}, at::TensorOptions().dtype(ctx->compute_type).device(base->device()));
-                    a.set_requires_grad(true);
-                    b.set_requires_grad(true);
-                    ctx->lora_a.push_back(std::move(a));
-                    ctx->lora_b.push_back(std::move(b));
-                    ctx->lora_names.push_back("lora_a_" + std::to_string(i) + "_" + std::to_string(k));
-                    ctx->lora_names.push_back("lora_b_" + std::to_string(i) + "_" + std::to_string(k));
+            double weighted_loss = 0.0;
+            int32_t completed_groups = 0;
+            for (const auto& [group_key, indexes] : groups) {
+                std::vector<int64_t> row_indexes;
+                row_indexes.reserve(indexes.size());
+                for (const size_t selected_index : indexes)
+                    row_indexes.push_back(
+                        static_cast<int64_t>(selected_index));
+
+                at::Tensor group_input;
+                at::Tensor group_targets;
+                at::Tensor group_attention;
+                void* group_input_ptr = input_ids_ptr;
+                void* group_target_ptr = target_mask_ptr;
+                void* group_attention_ptr = attention_mask_ptr;
+                bool local_prepared = true;
+                try {
+                    if (input_batch != 1) {
+                        auto row_tensor = at::tensor(
+                            row_indexes,
+                            input_ids.options().dtype(at::kLong));
+                        group_input = input_ids.index_select(0, row_tensor);
+                        group_targets = target_mask.index_select(0, row_tensor);
+                        group_input_ptr = &group_input;
+                        group_target_ptr = &group_targets;
+                        if (attention_mask.defined()) {
+                            group_attention = attention_mask.index_select(
+                                0, row_tensor);
+                            group_attention_ptr = &group_attention;
+                        }
+                    }
+                } catch (...) {
+                    local_prepared = false;
                 }
-            } else {
-                // Linear attention: in_proj_qkv, in_proj_z, out_proj
-                int64_t proj_indices[] = {2, 3, 10};  // qkv, z, out
-                for (int k = 0; k < 3; k++) {
-                    auto* base = ctx->weight_ptrs[w_offset + proj_indices[k]];
-                    int64_t out_f = base->size(0), in_f = base->size(1);
-                    auto a = at::randn({lora_rank, in_f}, at::TensorOptions().dtype(ctx->compute_type).device(base->device())) * 0.01;
-                    auto b = at::zeros({out_f, lora_rank}, at::TensorOptions().dtype(ctx->compute_type).device(base->device()));
-                    a.set_requires_grad(true);
-                    b.set_requires_grad(true);
-                    ctx->lora_a.push_back(std::move(a));
-                    ctx->lora_b.push_back(std::move(b));
-                    ctx->lora_names.push_back("lora_a_" + std::to_string(i) + "_" + std::to_string(k));
-                    ctx->lora_names.push_back("lora_b_" + std::to_string(i) + "_" + std::to_string(k));
+                TORCH_CHECK(adapter_collective_all_succeeded(
+                        ctx, local_prepared),
+                    "heterogeneous adapter group preparation failed on at "
+                    "least one distributed rank: ", group_key);
+
+                std::vector<TrainingContext::LoRAAdapter> group;
+                group.reserve(indexes.size());
+                for (const size_t selected_index : indexes)
+                    group.push_back(std::move(selected[selected_index]));
+                ctx->adapters.swap(group);
+                const int32_t group_size =
+                    static_cast<int32_t>(indexes.size());
+                const int32_t group_rank = static_cast<int32_t>(
+                    ctx->adapters.front().rank);
+                at::Tensor group_loss_numerators;
+                at::Tensor group_token_counts;
+                const double group_loss = qwen36_train_multi_lora_impl(
+                    ctx, group_input_ptr, group_target_ptr,
+                    group_attention_ptr, group_size, group_rank,
+                    DynamicMultiLoraMode::TrainOnly, nullptr,
+                    report_requested ? &group_loss_numerators : nullptr,
+                    report_requested ? &group_token_counts : nullptr);
+
+                group.swap(ctx->adapters);
+                for (size_t group_index = 0;
+                     group_index < indexes.size(); ++group_index) {
+                    selected[indexes[group_index]] =
+                        std::move(group[group_index]);
+                }
+                ++completed_groups;
+                bool local_group_succeeded =
+                    std::isfinite(group_loss) && group_loss >= 0.0;
+                const char* fail_after = std::getenv(
+                    "QWEN36_TEST_FAIL_HETERO_GROUP_AFTER");
+                if (fail_after && fail_after[0] != '\0') {
+                    char* end = nullptr;
+                    const long requested = std::strtol(
+                        fail_after, &end, 10);
+                    local_group_succeeded = local_group_succeeded &&
+                        end && *end == '\0' && requested > 0 &&
+                        completed_groups != requested;
                 }
+                TORCH_CHECK(adapter_collective_all_succeeded(
+                        ctx, local_group_succeeded),
+                    "heterogeneous adapter group failed on at least one "
+                    "distributed rank: ", group_key);
+                if (report_requested) {
+                    TORCH_CHECK(group_loss_numerators.defined() &&
+                            group_token_counts.defined() &&
+                            group_loss_numerators.numel() == group_size &&
+                            group_token_counts.numel() == group_size,
+                        "heterogeneous adapter group loss report shape mismatch: ",
+                        group_key);
+                    auto group_indexes = at::tensor(
+                        row_indexes,
+                        input_ids.options().dtype(at::kLong));
+                    selected_loss_numerators.index_copy_(
+                        0, group_indexes, group_loss_numerators);
+                    selected_token_counts.index_copy_(
+                        0, group_indexes, group_token_counts);
+                }
+                weighted_loss += group_loss * indexes.size();
             }
-            offset += lora_count;
+            train_loss = weighted_loss / static_cast<double>(n_adapters);
         }
 
-        // Initialize Adam state (FP32 for numerical stability, even if params are BF16)
-        for (size_t i = 0; i < ctx->lora_a.size(); i++) {
-            auto opts_f32 = at::TensorOptions().dtype(at::kFloat).device(ctx->lora_a[i].device());
-            ctx->adam_m.push_back(at::zeros(ctx->lora_a[i].sizes(), opts_f32));
-            ctx->adam_m.push_back(at::zeros(ctx->lora_b[i].sizes(), opts_f32));
-            ctx->adam_v.push_back(at::zeros(ctx->lora_a[i].sizes(), opts_f32));
-            ctx->adam_v.push_back(at::zeros(ctx->lora_b[i].sizes(), opts_f32));
+        if (report_requested) {
+            global_adapter_loss_values = global_dynamic_adapter_losses(
+                ctx, selected_loss_numerators, selected_token_counts);
+            auto losses_cpu = global_adapter_loss_values.to(
+                at::TensorOptions().device(at::kCPU).dtype(at::kDouble));
+            const auto* losses = losses_cpu.data_ptr<double>();
+            double reported_loss = 0.0;
+            for (int32_t index = 0; index < n_adapters; ++index) {
+                adapter_losses_out[index] = losses[index];
+                reported_loss += losses[index];
+            }
+            train_loss = reported_loss / static_cast<double>(n_adapters);
+            *aggregate_loss_out = train_loss;
         }
 
-        fprintf(stderr, "[q36_ctx] created: %ld layers, %ld LoRA params, %ld Adam states\n",
-            (long)num_layers, (long)ctx->lora_a.size(), (long)ctx->adam_m.size());
-        return ctx;
+        ctx->adapters.swap(selected);
+        selected_registry_installed = true;
+        int32_t finalizer_phase = 0;
+        const double finalizer_result = qwen36_train_multi_lora_impl(
+            ctx, input_ids_ptr, target_mask_ptr, attention_mask_ptr,
+            n_adapters, 0, DynamicMultiLoraMode::FinalizeOnly,
+            &finalizer_phase, nullptr, nullptr, false, nullptr, false, true);
+        selected.swap(ctx->adapters);
+        selected_registry_installed = false;
+        if (finalizer_phase < 1) {
+            // Match peers already waiting in the token preflight phase when
+            // this rank failed during finalizer-local preparation.
+            adapter_collective_all_succeeded(ctx, false);
+        }
+        const bool local_finalizer_succeeded =
+            std::isfinite(finalizer_result) && finalizer_result >= 0.0;
+        TORCH_CHECK(adapter_collective_all_succeeded(
+                ctx, local_finalizer_succeeded),
+            "heterogeneous adapter finalizer failed on at least one "
+            "distributed rank");
+        restore_registry();
+        release_dynamic_adam_shadows(ctx);
+        if (ctx) ctx->dynamic_adam_transaction_active = false;
+        page_cold_dynamic_adam_state(ctx, adapter_ids, n_adapters);
+        return train_loss;
     } catch (const std::exception& e) {
-        fprintf(stderr, "[q36] create FAILED: %s\n", e.what());
-        return nullptr;
+        recover();
+        fprintf(stderr, "[train_multi_selected_v2] FAILED: %s\n", e.what());
+        return -1.0;
+    } catch (...) {
+        recover();
+        fprintf(stderr,
+            "[train_multi_selected_v2] FAILED: unknown exception\n");
+        return -1.0;
     }
 }
 
-// Set MTP weights on an existing training context.
-// Called after create_training_context if MTP is enabled.
-__attribute__((visibility("default"))) void qwen36_set_mtp_weights(
+__attribute__((visibility("default"))) double
+qwen36_train_multi_lora_selected_v2(
     void* ctx_ptr,
-    void* mtp_fc_ptr,
-    void* mtp_pre_fc_norm_emb_ptr,
-    void* mtp_pre_fc_norm_hidden_ptr,
-    void* mtp_norm_ptr,
-    void** mtp_layer_weight_ptrs, int64_t num_mtp_layer_weights,
-    void* mtp_layer_configs_ptr, int64_t num_mtp_layers
+    void* input_ids_ptr,
+    void* target_mask_ptr,
+    void* attention_mask_ptr,
+    const int64_t* adapter_ids,
+    int32_t n_adapters
 ) {
-    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-    ctx->has_mtp = true;
-    ctx->mtp_fc = reinterpret_cast<at::Tensor*>(mtp_fc_ptr);
-    ctx->mtp_pre_fc_norm_emb = reinterpret_cast<at::Tensor*>(mtp_pre_fc_norm_emb_ptr);
-    ctx->mtp_pre_fc_norm_hidden = reinterpret_cast<at::Tensor*>(mtp_pre_fc_norm_hidden_ptr);
-    ctx->mtp_norm = reinterpret_cast<at::Tensor*>(mtp_norm_ptr);
+    return qwen36_train_multi_lora_selected_impl(
+        ctx_ptr, input_ids_ptr, target_mask_ptr, attention_mask_ptr,
+        adapter_ids, n_adapters, nullptr, nullptr, 0);
+}
 
-    auto** wp = reinterpret_cast<at::Tensor**>(mtp_layer_weight_ptrs);
-    for (int64_t i = 0; i < num_mtp_layer_weights; i++) {
-        ctx->mtp_layer_weights.push_back(wp[i]);
-    }
+__attribute__((visibility("default"))) int32_t
+qwen36_train_multi_lora_selected_v3(
+    void* ctx_ptr,
+    void* input_ids_ptr,
+    void* target_mask_ptr,
+    void* attention_mask_ptr,
+    const int64_t* adapter_ids,
+    int32_t n_adapters,
+    double* aggregate_loss_out,
+    double* adapter_losses_out,
+    int32_t adapter_loss_capacity
+) {
+    const double loss = qwen36_train_multi_lora_selected_impl(
+        ctx_ptr, input_ids_ptr, target_mask_ptr, attention_mask_ptr,
+        adapter_ids, n_adapters, aggregate_loss_out, adapter_losses_out,
+        adapter_loss_capacity);
+    return std::isfinite(loss) && loss >= 0.0 ? 0 : -1;
+}
 
-    auto* lcfgs = reinterpret_cast<LayerConfig*>(mtp_layer_configs_ptr);
-    for (int64_t i = 0; i < num_mtp_layers; i++) {
-        ctx->mtp_layer_configs.push_back(lcfgs[i]);
-    }
+// Set NCCL communicator for Expert Parallel all-reduce
+// Creates NCCL communicator directly in C++ using env vars RANK/WORLD_SIZE.
+// Rank 0 generates unique ID and writes to /tmp/rustrain-nccl/nccl-id.bin
+// Other ranks read it. All ranks call ncclCommInitRank.
+// Returns 0 on success, -1 on failure.
+// Process-level NCCL singleton — created once, reused across sessions.
+static ncclComm_t g_nccl_comm = nullptr;
+static cudaStream_t g_nccl_stream = nullptr;
+static ncclComm_t g_tp_comm = nullptr;
+static cudaStream_t g_tp_stream = nullptr;
+static ncclComm_t g_cp_comm = nullptr;
+static cudaStream_t g_cp_stream = nullptr;
+static ncclComm_t g_ep_comm = nullptr;
+static cudaStream_t g_ep_stream = nullptr;
+static ncclComm_t g_dp_comm = nullptr;
+static cudaStream_t g_dp_stream = nullptr;
+static ncclComm_t g_pp_comm = nullptr;
+static cudaStream_t g_pp_stream = nullptr;
+static ncclComm_t g_pp_control_comm = nullptr;
+static bool g_nccl_initialized = false;
+static bool g_nccl_cleanup_registered = false;
+static int g_parallel_rank = 0;
+static int g_parallel_world_size = 1;
+static int g_parallel_tp_rank = 0;
+static int g_parallel_tp_size = 1;
+static int g_parallel_tp_color = 0;
+static int g_parallel_cp_rank = 0;
+static int g_parallel_cp_size = 1;
+static int g_parallel_cp_color = 0;
+static int g_parallel_ep_rank = 0;
+static int g_parallel_ep_size = 1;
+static int g_parallel_ep_color = 0;
+static int g_parallel_dp_rank = 0;
+static int g_parallel_dp_size = 1;
+static int g_parallel_dp_color = 0;
+static int g_parallel_pp_rank = 0;
+static int g_parallel_pp_size = 1;
+static int g_parallel_pp_color = 0;
+
+static void qwen36_destroy_process_communicators() {
+    if (g_pp_control_comm) ncclCommDestroy(g_pp_control_comm);
+    if (g_pp_comm) ncclCommDestroy(g_pp_comm);
+    if (g_dp_comm) ncclCommDestroy(g_dp_comm);
+    if (g_ep_comm) ncclCommDestroy(g_ep_comm);
+    if (g_cp_comm) ncclCommDestroy(g_cp_comm);
+    if (g_tp_comm) ncclCommDestroy(g_tp_comm);
+    if (g_nccl_comm) ncclCommDestroy(g_nccl_comm);
+    g_pp_comm = nullptr;
+    g_pp_control_comm = nullptr;
+    g_dp_comm = nullptr;
+    g_ep_comm = nullptr;
+    g_cp_comm = nullptr;
+    g_tp_comm = nullptr;
+    g_nccl_comm = nullptr;
+    g_nccl_initialized = false;
+}
 
-    fprintf(stderr, "[q36_ctx] MTP set: %ld MTP layers, %ld MTP weight pointers\n",
-        (long)num_mtp_layers, (long)num_mtp_layer_weights);
+static bool same_cached_parallel_topology(
+    int rank, int world_size,
+    int tp_rank, int tp_size, int tp_color,
+    int cp_rank, int cp_size, int cp_color,
+    int ep_rank, int ep_size, int ep_color,
+    int dp_rank, int dp_size, int dp_color,
+    int pp_rank, int pp_size, int pp_color
+) {
+    return g_parallel_rank == rank &&
+        g_parallel_world_size == world_size &&
+        g_parallel_tp_rank == tp_rank &&
+        g_parallel_tp_size == tp_size &&
+        g_parallel_tp_color == tp_color &&
+        g_parallel_cp_rank == cp_rank &&
+        g_parallel_cp_size == cp_size &&
+        g_parallel_cp_color == cp_color &&
+        g_parallel_ep_rank == ep_rank &&
+        g_parallel_ep_size == ep_size &&
+        g_parallel_ep_color == ep_color &&
+        g_parallel_dp_rank == dp_rank &&
+        g_parallel_dp_size == dp_size &&
+        g_parallel_dp_color == dp_color &&
+        g_parallel_pp_rank == pp_rank &&
+        g_parallel_pp_size == pp_size &&
+        g_parallel_pp_color == pp_color;
 }
+static int g_cuda_device = 0;
 
-// Single training step: forward + loss + backward + Adam update
-// Returns loss value, or -1 on error.
-__attribute__((visibility("default"))) double qwen36_train_step(
+// Set CUDA device — called from Rust worker before any GPU operation.
+// Ensures PyTorch initializes CUDA context on the correct device.
+__attribute__((visibility("default"))) void qwen36_set_cuda_device(int32_t device) {
+    // Use PyTorch's device API — this updates both cudaSetDevice AND
+    // PyTorch's internal device tracking (c10::cuda::current_device).
+    // Must be called before any GPU operation in exec'd worker processes.
+    c10::cuda::set_device(device);
+    cudaSetDevice(device);
+    g_cuda_device = device;
+    // Force PyTorch to create CUDA context on this device
+    auto opts = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, device);
+    auto dummy = at::empty({1}, opts);
+    dummy.sizes();  // touch to ensure materialization
+}
+
+static int32_t qwen36_init_parallel_nccl_impl(
     void* ctx_ptr,
-    void* input_ids_ptr,
-    void* target_mask_ptr,
-    void* attention_mask_ptr
+    int rank, int world_size,
+    int tp_rank, int tp_size, int tp_color,
+    int cp_rank, int cp_size, int cp_color,
+    int ep_rank, int ep_size, int ep_color,
+    int dp_rank, int dp_size, int dp_color,
+    int pp_rank, int pp_size, int pp_color,
+    bool synchronize_parameters
 ) {
-    try {
-        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-        // Set CUDA device for EP
-        if (ctx->nccl_comm) {
-            c10::cuda::set_device(ctx->ep_rank);
-            cudaSetDevice(ctx->ep_rank);
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+
+    if (!ctx) return -1;
+
+    const int64_t topology_product = static_cast<int64_t>(tp_size) * cp_size *
+        ep_size * dp_size * pp_size;
+    if (rank < 0 || rank >= world_size || world_size <= 0 ||
+        tp_rank < 0 || tp_rank >= tp_size || tp_size <= 0 || tp_color < 0 ||
+        cp_rank < 0 || cp_rank >= cp_size || cp_size <= 0 || cp_color < 0 ||
+        ep_rank < 0 || ep_rank >= ep_size || ep_size <= 0 || ep_color < 0 ||
+        dp_rank < 0 || dp_rank >= dp_size || dp_size <= 0 || dp_color < 0 ||
+        pp_rank < 0 || pp_rank >= pp_size || pp_size <= 0 || pp_color < 0 ||
+        topology_product > std::numeric_limits<int>::max() ||
+        world_size != topology_product) {
+        ctx->topology_invalid = true;
+        fprintf(stderr,
+            "[parallel_nccl] invalid topology: rank=%d world=%d "
+            "tp=%d/%d color=%d cp=%d/%d color=%d ep=%d/%d color=%d "
+            "dp=%d/%d color=%d pp=%d/%d color=%d\n",
+            rank, world_size, tp_rank, tp_size, tp_color,
+            cp_rank, cp_size, cp_color, ep_rank, ep_size, ep_color,
+            dp_rank, dp_size, dp_color, pp_rank, pp_size, pp_color);
+        return -1;
+    }
+    if (ctx->tp_world_size != tp_size || ctx->tp_rank != tp_rank ||
+        ctx->cp_world_size != cp_size || ctx->cp_rank != cp_rank ||
+        ctx->ep_world_size != ep_size || ctx->dp_world_size != dp_size ||
+        ctx->ep_rank != ep_rank || ctx->dp_rank != dp_rank ||
+        ctx->pp_world_size != pp_size || ctx->pp_rank != pp_rank ||
+        ctx->expert_parallel != (ep_size > 1) ||
+        ctx->data_parallel != (dp_size > 1)) {
+        ctx->topology_invalid = true;
+        fprintf(stderr,
+            "[parallel_nccl] init topology does not match context: "
+            "context tp=%d/%d cp=%d/%d ep=%d/%d dp=%d/%d pp=%d/%d, "
+            "init tp=%d/%d cp=%d/%d ep=%d/%d dp=%d/%d pp=%d/%d\n",
+            ctx->tp_rank, ctx->tp_world_size, ctx->cp_rank, ctx->cp_world_size,
+            ctx->ep_rank, ctx->ep_world_size, ctx->dp_rank, ctx->dp_world_size,
+            ctx->pp_rank, ctx->pp_world_size,
+            tp_rank, tp_size, cp_rank, cp_size, ep_rank, ep_size,
+            dp_rank, dp_size, pp_rank, pp_size);
+        return -1;
+    }
+    ctx->topology_invalid = false;
+
+    // If already initialized, just set the pointer on this context
+    if (g_nccl_initialized) {
+        if (!same_cached_parallel_topology(
+                rank, world_size, tp_rank, tp_size, tp_color,
+                cp_rank, cp_size, cp_color,
+                ep_rank, ep_size, ep_color,
+                dp_rank, dp_size, dp_color,
+                pp_rank, pp_size, pp_color)) {
+            ctx->topology_invalid = true;
+            fprintf(stderr,
+                "[parallel_nccl] process communicator topology cannot change after initialization\n");
+            return -1;
         }
-        auto& input_ids = *reinterpret_cast<at::Tensor*>(input_ids_ptr);
-        auto& target_mask = *reinterpret_cast<at::Tensor*>(target_mask_ptr);
-        if (attention_mask_ptr) {
-            ctx->attention_mask = *reinterpret_cast<at::Tensor*>(attention_mask_ptr);
+        ctx->expert_parallel = ep_size > 1;
+        ctx->ep_rank = ep_rank;
+        ctx->ep_world_size = ep_size;
+        ctx->data_parallel = dp_size > 1;
+        ctx->dp_rank = dp_rank;
+        ctx->dp_world_size = dp_size;
+        const char* local_rank_str2 = getenv("LOCAL_RANK");
+        ctx->cuda_device = local_rank_str2 ? atoi(local_rank_str2) : g_cuda_device;
+        ctx->tp_world_size = tp_size;
+        ctx->tp_rank = tp_rank;
+        ctx->cp_world_size = cp_size;
+        ctx->cp_rank = cp_rank;
+        ctx->pp_world_size = pp_size;
+        ctx->pp_rank = pp_rank;
+        ctx->topology_invalid = false;
+        ctx->tp_comm = tp_size > 1 ? g_tp_comm : nullptr;
+        ctx->tp_stream = tp_size > 1 ? g_tp_stream : nullptr;
+        ctx->cp_comm = cp_size > 1 ? g_cp_comm : nullptr;
+        ctx->cp_stream = cp_size > 1 ? g_cp_stream : nullptr;
+        ctx->nccl_comm = ep_size > 1 ? g_ep_comm : nullptr;
+        ctx->nccl_stream = ep_size > 1 ? g_ep_stream : nullptr;
+        ctx->dp_comm = dp_size > 1 ? g_dp_comm : nullptr;
+        ctx->dp_stream = dp_size > 1 ? g_dp_stream : nullptr;
+        ctx->pp_comm = pp_size > 1 ? g_pp_comm : nullptr;
+        ctx->pp_stream = pp_size > 1 ? g_pp_stream : nullptr;
+        ctx->pp_control_comm = pp_size > 1 ? g_pp_control_comm : nullptr;
+        void* layer_comm = ep_size > 1 ? (void*)g_ep_comm : nullptr;
+        void* layer_stream = ep_size > 1 ? (void*)g_ep_stream : nullptr;
+        for (auto& lc : ctx->layer_configs) {
+            lc.nccl_comm = layer_comm;
+            lc.nccl_stream = layer_stream;
+        }
+        for (auto& lc : ctx->mtp_layer_configs) {
+            lc.nccl_comm = layer_comm;
+            lc.nccl_stream = layer_stream;
         }
+        if (synchronize_parameters) {
+            validate_fixed_collective_registry(ctx);
+            synchronize_fixed_replicated_lora_parameters(ctx);
+        }
+        return 0;
+    }
+    if (world_size <= 1) return 0;  // no EP needed
 
-        // Forward: checkpoint (default) or fused layer (QWEN36_FUSED_LAYER=1)
-        bool use_fused = getenv("QWEN36_FUSED_LAYER");
-        auto hidden = use_fused
-            ? forward_full_fused(ctx, input_ids)
-            : ctx->use_checkpoint
-                ? forward_full_checkpoint(ctx, input_ids)
-                : forward_full(ctx, input_ids);
+    // Set CUDA device and initialize PyTorch CUDA context on this device.
+    // NCCL communicator binds to the current CUDA context. If PyTorch hasn't
+    // initialized CUDA on this device yet, NCCL gets a wrong context → "invalid argument".
+    // Creating a dummy tensor forces PyTorch to initialize its CUDA context.
+    const char* local_rank_str = getenv("LOCAL_RANK");
+    int local_rank = local_rank_str ? atoi(local_rank_str) : rank;
+    c10::cuda::set_device(local_rank);
+    cudaSetDevice(local_rank);
+    g_cuda_device = local_rank;
+    ctx->cuda_device = local_rank;
+    {
+        // Force PyTorch CUDA context initialization on this device
+        auto opts = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, local_rank);
+        auto dummy = at::empty({1}, opts);
+        dummy.sizes();  // touch to ensure materialization
+    }
+
+    // Exchange unique ID via file
+    // Rank 0 generates ID, writes to file, then writes "ready" sentinel.
+    // Other ranks wait for "ready" file, then read ID.
+    // This ensures rank 0's write is visible before others read.
+    std::string rendezvous_dir;
+    if (!nccl_sync_dir(&rendezvous_dir)) {
+        ctx->topology_invalid = true;
+        return -1;
+    }
+    const std::string id_path = rendezvous_dir + "/nccl-id.bin";
+    const std::string ready_path = rendezvous_dir + "/nccl-ready.txt";
+    ncclUniqueId unique_id;
+    if (rank == 0) {
+        // Clean up old files first
+        remove(ready_path.c_str());
+        for (int peer = 0; peer < world_size; ++peer) {
+            const std::string stale_barrier =
+                rendezvous_dir + "/barrier_" + std::to_string(peer);
+            remove(stale_barrier.c_str());
+        }
+        if (ncclGetUniqueId(&unique_id) != ncclSuccess) {
+            fprintf(stderr, "[parallel_nccl] failed to generate NCCL unique ID\n");
+            ctx->topology_invalid = true;
+            return -1;
+        }
+        const std::string id_temporary = id_path + ".tmp." +
+            std::to_string(static_cast<long long>(getpid()));
+        FILE* f = fopen(id_temporary.c_str(), "wb");
+        const bool id_written = f &&
+            fwrite(&unique_id, sizeof(unique_id), 1, f) == 1;
+        const bool id_closed = f && fclose(f) == 0;
+        f = nullptr;
+        if (!id_written || !id_closed ||
+            rename(id_temporary.c_str(), id_path.c_str()) != 0) {
+            remove(id_temporary.c_str());
+            fprintf(stderr, "[parallel_nccl] failed to publish NCCL unique ID\n");
+            ctx->topology_invalid = true;
+            return -1;
+        }
+        // Publish the ready sentinel only after the ID rename is visible.
+        const std::string ready_temporary = ready_path + ".tmp." +
+            std::to_string(static_cast<long long>(getpid()));
+        FILE* rf = fopen(ready_temporary.c_str(), "w");
+        const bool ready_written = rf && fprintf(rf, "ready\n") >= 0;
+        const bool ready_closed = rf && fclose(rf) == 0;
+        rf = nullptr;
+        if (!ready_written || !ready_closed ||
+            rename(ready_temporary.c_str(), ready_path.c_str()) != 0) {
+            remove(ready_temporary.c_str());
+            fprintf(stderr, "[parallel_nccl] failed to publish ready sentinel\n");
+            ctx->topology_invalid = true;
+            return -1;
+        }
+    } else {
+        // Wait for ready sentinel
+        bool ready = false;
+        for (int i = 0; i < 600; i++) {
+            FILE* rf = fopen(ready_path.c_str(), "r");
+            if (rf) { fclose(rf); ready = true; break; }
+            usleep(10000);  // 10ms
+        }
+        if (!ready) {
+            fprintf(stderr,
+                "[parallel_nccl] rank %d timed out waiting for ready sentinel\n",
+                rank);
+            ctx->topology_invalid = true;
+            return -1;
+        }
+        // Now read ID file
+        FILE* f = fopen(id_path.c_str(), "rb");
+        if (!f || fread(&unique_id, sizeof(unique_id), 1, f) != 1) {
+            fprintf(stderr, "[parallel_nccl] rank %d: failed to read ID file\n", rank);
+            if (f) fclose(f);
+            ctx->topology_invalid = true;
+            return -1;
+        }
+        fclose(f);
+    }
 
-        // Debug: GPU memory after forward
-        {
-            size_t free, total;
-            cudaMemGetInfo(&free, &total);
+    // Barrier: ensure all ranks reach ncclCommInitRank simultaneously.
+    // Without this, rank 0 (fast load_model) reaches ncclCommInitRank before
+    // rank 3 (slow load_model) → NCCL timeout.
+    {
+        const char* barrier_dir = rendezvous_dir.c_str();
+        char bpath[256];
+        snprintf(bpath, sizeof(bpath), "%s/barrier_%d", barrier_dir, rank);
+        FILE* bf = fopen(bpath, "w");
+        const bool barrier_written = bf && fprintf(bf, "1\n") >= 0;
+        const bool barrier_closed = bf && fclose(bf) == 0;
+        bf = nullptr;
+        if (!barrier_written || !barrier_closed) {
+            fprintf(stderr,
+                "[parallel_nccl] rank %d failed to publish barrier file\n", rank);
+            ctx->topology_invalid = true;
+            return -1;
+        }
+        // Wait for all ranks
+        for (int i = 0; i < world_size; i++) {
+            char p[256];
+            snprintf(p, sizeof(p), "%s/barrier_%d", barrier_dir, i);
+            bool peer_ready = false;
+            for (int w = 0; w < 6000; w++) {  // 60s timeout
+                FILE* f2 = fopen(p, "r");
+                if (f2) { fclose(f2); peer_ready = true; break; }
+                usleep(10000);
+            }
+            if (!peer_ready) {
+                fprintf(stderr,
+                    "[parallel_nccl] rank %d timed out waiting for rank %d barrier\n",
+                    rank, i);
+                ctx->topology_invalid = true;
+                return -1;
+            }
         }
+    }
 
-        // Main loss — compute_loss does chunked CE with immediate backward per chunk.
-        // This avoids accumulating 250 chunks of [512, vocab] logits in autograd graph.
-        auto loss = compute_loss(ctx, hidden, input_ids, target_mask, ctx->vocab_size);
+    // Initialize communicator — only once per process
+    ncclComm_t comm;
+    ncclResult_t err = ncclCommInitRank(&comm, world_size, unique_id, rank);
+    if (err != ncclSuccess) {
+        fprintf(stderr, "[ep_nccl] ncclCommInitRank failed: %d (%s)\n", err, ncclGetErrorString(err));
+        ctx->topology_invalid = true;
+        return -1;
+    }
 
-        // compute_loss did chunked CE backward on detached hidden,
-        // accumulated gradient into hidden.grad().
-        // Now trigger main model backward using hidden's gradient.
-        double loss_val = loss.item<double>();
+    cudaStream_t nccl_stream = nullptr;
 
-        // Debug: GPU memory after CE backward
-        {
-            size_t free, total;
-            cudaMemGetInfo(&free, &total);
+    // A divergent axis size would make ranks execute different split counts
+    // and hang. Establish size consensus on the world communicator first.
+    try {
+        auto topology_sizes = at::tensor(
+            std::vector<int32_t>{tp_size, cp_size, ep_size, dp_size, pp_size},
+            at::TensorOptions().device(at::kCUDA, local_rank).dtype(at::kInt));
+        auto minimum_sizes = topology_sizes.clone();
+        auto maximum_sizes = topology_sizes.clone();
+        auto topology_stream =
+            c10::cuda::getCurrentCUDAStream(local_rank).stream();
+        const ncclResult_t min_error = ncclAllReduce(
+            minimum_sizes.data_ptr<int32_t>(), minimum_sizes.data_ptr<int32_t>(),
+            5, ncclInt32, ncclMin, comm, topology_stream);
+        const ncclResult_t max_error = ncclAllReduce(
+            maximum_sizes.data_ptr<int32_t>(), maximum_sizes.data_ptr<int32_t>(),
+            5, ncclInt32, ncclMax, comm, topology_stream);
+        if (min_error != ncclSuccess || max_error != ncclSuccess) {
+            fprintf(stderr,
+                "[parallel_nccl] topology size consensus failed: min=%s max=%s\n",
+                ncclGetErrorString(min_error), ncclGetErrorString(max_error));
+            ncclCommDestroy(comm);
+            ctx->topology_invalid = true;
+            return -1;
         }
-
-        // Release CE's retained graph (hidden_normed etc.) before backward.
-        // CE gradient is already accumulated into hidden.grad().
-        // hidden.detach() breaks the autograd graph from CE.
-        hidden = hidden.detach();
-        if (hidden.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache();
-
-        // Debug: after releasing CE graph
-        {
-            size_t free, total;
-            cudaMemGetInfo(&free, &total);
+        const auto minimum_cpu = minimum_sizes.to(at::kCPU);
+        const auto maximum_cpu = maximum_sizes.to(at::kCPU);
+        bool sizes_match = true;
+        for (int axis = 0; axis < 5; ++axis) {
+            sizes_match = sizes_match &&
+                minimum_cpu.data_ptr<int32_t>()[axis] ==
+                    maximum_cpu.data_ptr<int32_t>()[axis];
         }
-
-        // Trigger main model backward.
-        // FusedLayer: PyTorch autograd handles backward per layer.
-        // Checkpoint: manual_group_backward recomputes each group with PyTorch.
-        if (use_fused && hidden.grad().defined()) {
-            hidden.backward(hidden.grad());
-        } else if (hidden.grad().defined() && !ctx->group_inputs.empty() && !getenv("QWEN36_SUBCKPT")) {
-            manual_group_backward(ctx, hidden.grad());
-        } else if (hidden.grad().defined()) {
-            hidden.backward(hidden.grad());
+        if (!sizes_match) {
+            fprintf(stderr,
+                "[parallel_nccl] TP/CP/EP/DP/PP sizes differ across ranks\n");
+            ncclCommDestroy(comm);
+            ctx->topology_invalid = true;
+            return -1;
         }
+    } catch (const std::exception& error) {
+        fprintf(stderr,
+            "[parallel_nccl] topology size consensus threw an exception: %s\n",
+            error.what());
+        ncclCommDestroy(comm);
+        ctx->topology_invalid = true;
+        return -1;
+    } catch (...) {
+        fprintf(stderr,
+            "[parallel_nccl] topology size consensus threw an unknown exception\n");
+        ncclCommDestroy(comm);
+        ctx->topology_invalid = true;
+        return -1;
+    }
 
-        // MTP loss (if enabled) — run AFTER main backward to reuse freed GPU memory
-        if (ctx->has_mtp && !getenv("QWEN36_DISABLE_MTP")) {
-            // Detach hidden so MTP forward doesn't rebuild main model graph
-            auto hidden_detached = hidden.detach().set_requires_grad(true);
-            auto mtp_hidden = mtp_forward(ctx, hidden_detached, input_ids);
-            auto mtp_loss = mtp_compute_loss(ctx, mtp_hidden, input_ids, target_mask);
-            // Backward MTP loss — frees MTP intermediate tensors immediately
-            mtp_loss.backward();
-            // Add MTP gradient to hidden's gradient (already populated by main backward)
-            if (hidden_detached.grad().defined()) {
-                if (hidden.grad().defined()) {
-                    hidden.grad().add_(hidden_detached.grad());
-                } else {
-                    hidden.mutable_grad() = hidden_detached.grad().clone();
-                }
-            }
+    // Build every sub-communicator before publishing process-global state.
+    // All ranks use the same TP, CP, EP, DP, PP split order.
+    ncclComm_t tp_comm = nullptr;
+    ncclComm_t cp_comm = nullptr;
+    ncclComm_t ep_comm = nullptr;
+    ncclComm_t dp_comm = nullptr;
+    ncclComm_t pp_comm = nullptr;
+    ncclComm_t pp_control_comm = nullptr;
+    auto destroy_local_communicators = [&]() {
+        if (pp_control_comm) ncclCommDestroy(pp_control_comm);
+        if (pp_comm) ncclCommDestroy(pp_comm);
+        if (dp_comm) ncclCommDestroy(dp_comm);
+        if (ep_comm) ncclCommDestroy(ep_comm);
+        if (cp_comm) ncclCommDestroy(cp_comm);
+        if (tp_comm) ncclCommDestroy(tp_comm);
+        ncclCommDestroy(comm);
+    };
+    auto split_axis = [&](const char* axis, int size, int color, int key,
+                          ncclComm_t* output) {
+        if (size <= 1) return true;
+        const ncclResult_t split_error = ncclCommSplit(
+            comm, color, key, output, nullptr);
+        if (split_error == ncclSuccess) return true;
+        fprintf(stderr, "[%s_nccl] ncclCommSplit failed: %d (%s)\n",
+            axis, split_error, ncclGetErrorString(split_error));
+        return false;
+    };
+    if (!split_axis("tp", tp_size, tp_color, tp_rank, &tp_comm) ||
+        !split_axis("cp", cp_size, cp_color, cp_rank, &cp_comm) ||
+        !split_axis("ep", ep_size, ep_color, ep_rank, &ep_comm) ||
+        !split_axis("dp", dp_size, dp_color, dp_rank, &dp_comm) ||
+        !split_axis("pp", pp_size, pp_color, pp_rank, &pp_comm)) {
+        destroy_local_communicators();
+        ctx->topology_invalid = true;
+        return -1;
+    }
+    if (pp_comm) {
+        const ncclResult_t control_error = ncclCommSplit(
+            pp_comm, 0, pp_rank, &pp_control_comm, nullptr);
+        if (control_error != ncclSuccess) {
+            fprintf(stderr, "[pp_nccl] control communicator split failed: %d (%s)\\n",
+                control_error, ncclGetErrorString(control_error));
+            destroy_local_communicators();
+            ctx->topology_invalid = true;
+            return -1;
         }
+    }
 
-        // ── Adam optimizer step — CUDA multi-tensor fused kernel ──
-        at::AutoGradMode guard(false);
-        ctx->step_count++;
-        ctx->lora_cache_valid = false;
-        ctx->lora_batch_valid = false;
-        double step_f = (double)ctx->step_count;
-        double bias_correction1 = 1.0 - std::pow(ctx->beta1, step_f);
-        double bias_correction2 = 1.0 - std::pow(ctx->beta2, step_f);
-        float lr_scaled = (float)(ctx->lr / bias_correction1);
-        float eps_scaled = (float)(ctx->eps / std::sqrt(bias_correction2));
-        float one_minus_b1 = (float)(1.0 - ctx->beta1);
-        float one_minus_b2 = (float)(1.0 - ctx->beta2);
-
-        // Collect all (param, grad, m, v, size) tuples from multi-LoRA + legacy
-        std::vector<void*> h_params, h_grads;
-        std::vector<float*> h_m, h_v;
-        std::vector<int> h_sizes;
-
-        // Multi-LoRA adapters
-        for (auto& adapter : ctx->adapters) {
-            for (auto& [layer_idx, pairs] : adapter.params) {
-                auto& adam_states = adapter.adam_state[layer_idx];
-                for (size_t i = 0; i < pairs.size(); i++) {
-                    auto& [a, b] = pairs[i];
-                    auto& [m_a, v_a, m_b, v_b] = adam_states[i];
-                    if (a.grad().defined() && a.scalar_type() == at::kBFloat16) {
-                        h_params.push_back(a.data_ptr());
-                        h_grads.push_back(a.grad().data_ptr());
-                        h_m.push_back((float*)m_a.data_ptr());
-                        h_v.push_back((float*)v_a.data_ptr());
-                        h_sizes.push_back((int)a.numel());
-                    }
-                    if (b.grad().defined() && b.scalar_type() == at::kBFloat16) {
-                        h_params.push_back(b.data_ptr());
-                        h_grads.push_back(b.grad().data_ptr());
-                        h_m.push_back((float*)m_b.data_ptr());
-                        h_v.push_back((float*)v_b.data_ptr());
-                        h_sizes.push_back((int)b.numel());
-                    }
-                }
-            }
-        }
-        // Legacy single-LoRA
-        size_t adam_idx = 0;
-        for (size_t i = 0; i < ctx->lora_a.size(); i++) {
-            {
-                auto& param = ctx->lora_a[i];
-                auto& grad = param.grad();
-                if (grad.defined() && param.scalar_type() == at::kBFloat16) {
-                    h_params.push_back(param.data_ptr());
-                    h_grads.push_back(grad.data_ptr());
-                    h_m.push_back((float*)ctx->adam_m[adam_idx].data_ptr());
-                    h_v.push_back((float*)ctx->adam_v[adam_idx].data_ptr());
-                    h_sizes.push_back((int)param.numel());
-                }
-            }
-            adam_idx++;
-            {
-                auto& param = ctx->lora_b[i];
-                auto& grad = param.grad();
-                if (grad.defined() && param.scalar_type() == at::kBFloat16) {
-                    h_params.push_back(param.data_ptr());
-                    h_grads.push_back(grad.data_ptr());
-                    h_m.push_back((float*)ctx->adam_m[adam_idx].data_ptr());
-                    h_v.push_back((float*)ctx->adam_v[adam_idx].data_ptr());
-                    h_sizes.push_back((int)param.numel());
-                }
-            }
-            adam_idx++;
-        }
-
-        int n_params = (int)h_params.size();
-        if (n_params > 0) {
-            // ── CUDA multi-tensor fused Adam: 1 launch for all params ──
-            // Ensure device buffers are large enough
-            ctx->adam_dev_bufs.ensure(n_params, ctx->lora_a.empty()
-                ? ctx->adapters[0].params.begin()->second[0].first
-                : ctx->lora_a[0]);
-
-            // Copy pointer arrays to device
-            auto opts_cpu_long = at::TensorOptions().dtype(at::kLong).device(at::kCPU);
-            auto opts_cpu_int  = at::TensorOptions().dtype(at::kInt).device(at::kCPU);
-            auto params_cpu = at::from_blob(h_params.data(), {n_params}, opts_cpu_long);
-            auto grads_cpu  = at::from_blob(h_grads.data(),  {n_params}, opts_cpu_long);
-            auto m_cpu      = at::from_blob(h_m.data(),      {n_params}, opts_cpu_long);
-            auto v_cpu      = at::from_blob(h_v.data(),      {n_params}, opts_cpu_long);
-            auto sizes_cpu  = at::from_blob(h_sizes.data(),  {n_params}, opts_cpu_int);
-            ctx->adam_dev_bufs.params_buf.copy_(params_cpu);
-            ctx->adam_dev_bufs.grads_buf.copy_(grads_cpu);
-            ctx->adam_dev_bufs.m_buf.copy_(m_cpu);
-            ctx->adam_dev_bufs.v_buf.copy_(v_cpu);
-            ctx->adam_dev_bufs.sizes_buf.copy_(sizes_cpu);
-
-            // Single kernel launch for ALL params
-            auto stream = c10::cuda::getCurrentCUDAStream().stream();
-            launch_fused_adam_multi(
-                (void**)ctx->adam_dev_bufs.params_buf.data_ptr(),
-                (void**)ctx->adam_dev_bufs.grads_buf.data_ptr(),
-                (float**)ctx->adam_dev_bufs.m_buf.data_ptr(),
-                (float**)ctx->adam_dev_bufs.v_buf.data_ptr(),
-                (int*)ctx->adam_dev_bufs.sizes_buf.data_ptr(),
-                n_params,
-                (float)ctx->beta1, (float)ctx->beta2,
-                lr_scaled, eps_scaled,
-                one_minus_b1, one_minus_b2,
-                (void*)stream
-            );
-        }
+    g_nccl_comm = comm;
+    g_nccl_stream = nccl_stream;
+    g_tp_comm = tp_comm;
+    g_tp_stream = nccl_stream;
+    g_cp_comm = cp_comm;
+    g_cp_stream = nccl_stream;
+    g_ep_comm = ep_comm;
+    g_ep_stream = nccl_stream;
+    g_dp_comm = dp_comm;
+    g_dp_stream = nccl_stream;
+    g_pp_comm = pp_comm;
+    g_pp_stream = nccl_stream;
+    g_pp_control_comm = pp_control_comm;
+    g_parallel_rank = rank;
+    g_parallel_world_size = world_size;
+    g_parallel_tp_rank = tp_rank;
+    g_parallel_tp_size = tp_size;
+    g_parallel_tp_color = tp_color;
+    g_parallel_cp_rank = cp_rank;
+    g_parallel_cp_size = cp_size;
+    g_parallel_cp_color = cp_color;
+    g_parallel_ep_rank = ep_rank;
+    g_parallel_ep_size = ep_size;
+    g_parallel_ep_color = ep_color;
+    g_parallel_dp_rank = dp_rank;
+    g_parallel_dp_size = dp_size;
+    g_parallel_dp_color = dp_color;
+    g_parallel_pp_rank = pp_rank;
+    g_parallel_pp_size = pp_size;
+    g_parallel_pp_color = pp_color;
+    g_nccl_initialized = true;
+    if (!g_nccl_cleanup_registered) {
+        std::atexit(qwen36_destroy_process_communicators);
+        g_nccl_cleanup_registered = true;
+    }
 
-        return loss_val;
-    } catch (const std::exception& e) {
-        fprintf(stderr, "[q36] train_step FAILED: %s\n", e.what());
-        return -1.0;
+    ctx->nccl_comm = ep_size > 1 ? g_ep_comm : nullptr;
+    ctx->nccl_stream = ep_size > 1 ? g_ep_stream : nullptr;
+    ctx->ep_rank = ep_rank;
+    ctx->ep_world_size = ep_size;
+    ctx->expert_parallel = ep_size > 1;
+    ctx->dp_comm = dp_size > 1 ? g_dp_comm : nullptr;
+    ctx->dp_stream = dp_size > 1 ? g_dp_stream : nullptr;
+    ctx->dp_rank = dp_rank;
+    ctx->dp_world_size = dp_size;
+    ctx->data_parallel = dp_size > 1;
+    ctx->tp_world_size = tp_size;
+    ctx->tp_rank = tp_rank;
+    ctx->tp_comm = tp_size > 1 ? g_tp_comm : nullptr;
+    ctx->tp_stream = tp_size > 1 ? g_tp_stream : nullptr;
+    ctx->cp_world_size = cp_size;
+    ctx->cp_rank = cp_rank;
+    ctx->cp_comm = cp_size > 1 ? g_cp_comm : nullptr;
+    ctx->cp_stream = cp_size > 1 ? g_cp_stream : nullptr;
+    ctx->pp_world_size = pp_size;
+    ctx->pp_rank = pp_rank;
+    ctx->pp_comm = pp_size > 1 ? g_pp_comm : nullptr;
+    ctx->pp_stream = pp_size > 1 ? g_pp_stream : nullptr;
+    ctx->pp_control_comm = pp_size > 1 ? g_pp_control_comm : nullptr;
+
+    // Propagate to layer configs
+    void* layer_comm = ep_size > 1 ? (void*)g_ep_comm : nullptr;
+    void* layer_stream = ep_size > 1 ? (void*)g_ep_stream : nullptr;
+    for (auto& lc : ctx->layer_configs) {
+        lc.nccl_comm = layer_comm;
+        lc.nccl_stream = layer_stream;
+    }
+    for (auto& lc : ctx->mtp_layer_configs) {
+        lc.nccl_comm = layer_comm;
+        lc.nccl_stream = layer_stream;
     }
-}
 
-// Get LoRA A tensor pointer by index
-__attribute__((visibility("default"))) void* qwen36_get_lora_a(void* ctx_ptr, int64_t index) {
-    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-    return &ctx->lora_a[index];
+    if (synchronize_parameters) {
+        validate_fixed_collective_registry(ctx);
+        synchronize_fixed_replicated_lora_parameters(ctx);
+    }
+    return 0;
 }
 
-// Get LoRA B tensor pointer by index
-__attribute__((visibility("default"))) void* qwen36_get_lora_b(void* ctx_ptr, int64_t index) {
+__attribute__((visibility("default"))) int32_t qwen36_init_parallel_nccl_v2(
+    void* ctx_ptr,
+    int32_t rank, int32_t world_size,
+    int32_t tp_rank, int32_t tp_size, int32_t tp_color,
+    int32_t cp_rank, int32_t cp_size, int32_t cp_color,
+    int32_t ep_rank, int32_t ep_size, int32_t ep_color,
+    int32_t dp_rank, int32_t dp_size, int32_t dp_color,
+    int32_t pp_rank, int32_t pp_size, int32_t pp_color
+) {
     auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-    return &ctx->lora_b[index];
-}
-
-// Free training context
-__attribute__((visibility("default"))) void qwen36_free_training_context(void* ctx_ptr) {
-    if (ctx_ptr) {
-        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-        // Don't destroy NCCL communicator — it's a process-level singleton
-        // (g_nccl_comm). It must survive context destruction so the next
-        // session can reuse it. Destroying it would break NCCL for all
-        // subsequent sessions in the same worker process.
-        // ncclCommDestroy is called only on process exit (via atexit or Drop).
-        delete ctx;
+    try {
+        return qwen36_init_parallel_nccl_impl(
+            ctx_ptr, rank, world_size,
+            tp_rank, tp_size, tp_color,
+            cp_rank, cp_size, cp_color,
+            ep_rank, ep_size, ep_color,
+            dp_rank, dp_size, dp_color,
+            pp_rank, pp_size, pp_color,
+            /*synchronize_parameters=*/true);
+    } catch (const std::exception& error) {
+        fprintf(stderr, "[parallel_nccl] initialization failed: %s\n", error.what());
+    } catch (...) {
+        fprintf(stderr, "[parallel_nccl] initialization failed with an unknown exception\n");
     }
+    if (ctx) ctx->topology_invalid = true;
+    return -1;
 }
 
-// ── Batched Multi-LoRA Training ──
-
-/// Compute max adapters that fit in available GPU memory.
-/// Based on per-adapter activation memory (dominant) + LoRA params + Adam state.
-static int64_t compute_n_max(
-    int64_t free_gpu_bytes, int64_t rank, int64_t seq,
-    int64_t hidden, int64_t group_size, int64_t num_layers
+// ABI25-compatible three-axis wrapper. ABI26 callers use the five-axis v2
+// entry point; retaining this symbol keeps native TP/EP/DP diagnostics simple.
+__attribute__((visibility("default"))) int32_t qwen36_init_parallel_nccl(
+    void* ctx_ptr,
+    int32_t rank, int32_t world_size,
+    int32_t tp_rank, int32_t tp_size, int32_t tp_color,
+    int32_t ep_rank, int32_t ep_size, int32_t ep_color,
+    int32_t dp_rank, int32_t dp_size, int32_t dp_color
 ) {
-    // Per-adapter activation memory (BF16, group_size=2 optimal):
-    // group_inputs: (num_layers / group_size) × seq × hidden × 2 bytes
-    // peak recompute: ~34 MB per layer pair × group_size (rough avg)
-    int64_t gs = group_size < 1 ? 1 : group_size;
-    int64_t num_groups = (num_layers + gs - 1) / gs;
-    int64_t group_input_mem = num_groups * seq * hidden * 2;  // BF16
-
-    // Average per-layer saved tensors during recompute:
-    // full_attn: ~17MB, linear_attn: ~25MB, moe: ~8MB → avg ~17MB
-    int64_t avg_layer_saved = 17 * 1024 * 1024;  // bytes
-    int64_t peak_mem = gs * avg_layer_saved;
-
-    // LoRA params + Adam state: 280 modules × (A+B) × (BF16 param + FP32 m + FP32 v)
-    // Per module: rank × hidden × 2 (BF16) + hidden × rank × 2 (BF16) + 2 × 4 (FP32 m+v)
-    // Simplified: 280 × rank × hidden × (2 + 2 + 8) = 280 × rank × hidden × 12
-    int64_t num_modules = 280;
-    int64_t lora_mem = num_modules * rank * hidden * 12;  // conservative
-
-    // CE peak: one chunk of [16384, vocab] logits at a time.
-    // BF16 logits (8GB) + FP32 CE loss (16GB) + grad (~8GB) ≈ 16GB.
-    // But backward releases immediately, so effective peak is lower.
-    int64_t ce_peak = 16384 * 248320 * 4;  // ~16GB (FP32 logits only, others freed by autograd)
-
-    // Attention intermediate: Q/K/V + attn_weights ≈ 4 × N × heads × seq × head_dim × 2 bytes
-    // Flash attention keeps this O(seq) not O(seq²), but still significant at seq=16K
-    int64_t attn_heads = 16;
-    int64_t head_dim = 256;
-    int64_t attn_mem = 4 * attn_heads * seq * head_dim * 2;  // per adapter, BF16
-
-    int64_t per_adapter = group_input_mem + peak_mem + lora_mem + attn_mem;
-    // Empirical multiplier: residual add, MoE routing, LoRA delta, etc.
-    // Multiplier scales with seq: small seq needs less (CPU dispatch dominant),
-    // large seq needs more (attention/MoE intermediates dominate).
-    // seq=512: 3x → n_max=100. seq=16K: 8x → n_max≈8 (auto-chunks N=20+).
-    int64_t mult = (seq > 4096) ? 8 : 3;
-    per_adapter = per_adapter * mult;
-    if (per_adapter <= 0) return 1;
-
-    // Reserve 15% for fragmentation + overhead, minus CE peak (constant overhead)
-    int64_t usable = (free_gpu_bytes - ce_peak) * 85 / 100;
-    if (usable < per_adapter) usable = free_gpu_bytes * 50 / 100;  // fallback: aggressive
-    int64_t n_max = usable / per_adapter;
-    return n_max < 1 ? 1 : n_max;
+    return qwen36_init_parallel_nccl_v2(
+        ctx_ptr, rank, world_size,
+        tp_rank, tp_size, tp_color,
+        0, 1, 0,
+        ep_rank, ep_size, ep_color,
+        dp_rank, dp_size, dp_color,
+        0, 1, 0);
 }
 
-/// Train all adapters in chunks. Each chunk: independent forward → loss → backward → Adam.
-/// Input is expanded to [N, seq] for each chunk.
-__attribute__((visibility("default"))) double qwen36_train_multi_lora(
+// Attach a shadow restore context to process-cached communicators without
+// broadcasting its temporary random LoRA initialization. Checkpoint tensors
+// replace every active parameter before the context can become live.
+__attribute__((visibility("default"))) int32_t
+qwen36_attach_parallel_nccl_no_sync_v2(
     void* ctx_ptr,
-    void* input_ids_ptr,
-    void* target_mask_ptr,
-    void* attention_mask_ptr,
-    int32_t n_total,
-    int32_t lora_rank
+    int32_t rank, int32_t world_size,
+    int32_t tp_rank, int32_t tp_size, int32_t tp_color,
+    int32_t cp_rank, int32_t cp_size, int32_t cp_color,
+    int32_t ep_rank, int32_t ep_size, int32_t ep_color,
+    int32_t dp_rank, int32_t dp_size, int32_t dp_color,
+    int32_t pp_rank, int32_t pp_size, int32_t pp_color
 ) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (world_size > 1 && !g_nccl_initialized) {
+        fprintf(stderr,
+            "[parallel_nccl] restore attach requires initialized process communicators\n");
+        if (ctx) ctx->topology_invalid = true;
+        return -1;
+    }
     try {
-        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-        if (ctx->nccl_comm) {
-            c10::cuda::set_device(ctx->ep_rank);
-            cudaSetDevice(ctx->ep_rank);
-        }
-
-        auto& input_ids = *reinterpret_cast<at::Tensor*>(input_ids_ptr);
-        auto& target_mask = *reinterpret_cast<at::Tensor*>(target_mask_ptr);
-        if (attention_mask_ptr) {
-            ctx->attention_mask = *reinterpret_cast<at::Tensor*>(attention_mask_ptr);
-        }
-
-        int64_t total_adapters = (int64_t)ctx->adapters.size();
-        if (total_adapters == 0) return -1.0;
-
-        // Compute N_max from available GPU memory.
-        // CRITICAL: all workers must agree on n_max to keep NCCL all-reduce in sync.
-        // Use file-based barrier: rank 0 computes n_max, writes to file, others read.
-        size_t free_mem, total_mem;
-        cudaMemGetInfo(&free_mem, &total_mem);
-        int64_t n_max;
-        if (ctx->nccl_comm && ctx->ep_world_size > 1) {
-            const char* sync_path = "/tmp/rustrain-nccl/nmax_sync.txt";
-            if (ctx->ep_rank == 0) {
-                n_max = compute_n_max(
-                    (int64_t)free_mem, lora_rank,
-                    input_ids.size(-1), 2048,
-                    ctx->group_size, ctx->num_layers
-                );
-                n_max = std::min(n_max, total_adapters);
-                if (n_max < 1) n_max = 1;
-                mkdir("/tmp/rustrain-nccl", 0777);
-                FILE* f = fopen(sync_path, "w");
-                fprintf(f, "%ld\n", (long)n_max);
-                fclose(f);
-            } else {
-                mkdir("/tmp/rustrain-nccl", 0777);
-                for (int i = 0; i < 600; i++) {
-                    FILE* f = fopen(sync_path, "r");
-                    if (f) { fscanf(f, "%ld", (long*)&n_max); fclose(f); break; }
-                    usleep(10000);
-                }
-            }
-        } else {
-            n_max = compute_n_max(
-                (int64_t)free_mem, lora_rank,
-                input_ids.size(-1), 2048,
-                ctx->group_size, ctx->num_layers
-            );
-        }
-        n_max = std::min(n_max, total_adapters);
-        if (n_max < 1) n_max = 1;
-
-        fprintf(stderr, "[train_multi] total=%ld n_max=%ld free=%.1fGB rank=%d\n",
-                (long)total_adapters, (long)n_max, (double)free_mem / 1e9, lora_rank);
-
-        double total_loss = 0.0;
-        int64_t num_chunks = (total_adapters + n_max - 1) / n_max;
+        return qwen36_init_parallel_nccl_impl(
+            ctx_ptr, rank, world_size,
+            tp_rank, tp_size, tp_color,
+            cp_rank, cp_size, cp_color,
+            ep_rank, ep_size, ep_color,
+            dp_rank, dp_size, dp_color,
+            pp_rank, pp_size, pp_color,
+            /*synchronize_parameters=*/false);
+    } catch (const std::exception& error) {
+        fprintf(stderr, "[parallel_nccl] restore attach failed: %s\n", error.what());
+    } catch (...) {
+        fprintf(stderr, "[parallel_nccl] restore attach failed with an unknown exception\n");
+    }
+    if (ctx) ctx->topology_invalid = true;
+    return -1;
+}
 
-        for (int64_t chunk = 0; chunk < num_chunks; chunk++) {
-            int64_t start = chunk * n_max;
-            int64_t end = std::min(start + n_max, total_adapters);
-            int64_t n = end - start;
+__attribute__((visibility("default"))) int32_t qwen36_attach_parallel_nccl_no_sync(
+    void* ctx_ptr,
+    int32_t rank, int32_t world_size,
+    int32_t tp_rank, int32_t tp_size, int32_t tp_color,
+    int32_t ep_rank, int32_t ep_size, int32_t ep_color,
+    int32_t dp_rank, int32_t dp_size, int32_t dp_color
+) {
+    return qwen36_attach_parallel_nccl_no_sync_v2(
+        ctx_ptr, rank, world_size,
+        tp_rank, tp_size, tp_color,
+        0, 1, 0,
+        ep_rank, ep_size, ep_color,
+        dp_rank, dp_size, dp_color,
+        0, 1, 0);
+}
 
-            // Invalidate cache for this chunk's adapter set
-            ctx->lora_batch_valid = false;
-            ctx->lora_cache_valid = false;
+__attribute__((visibility("default"))) int32_t qwen36_init_nccl(
+    void* ctx_ptr
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx) return -1;
+    const char* rank_str = getenv("RANK");
+    const char* world_str = getenv("WORLD_SIZE");
+    if (!rank_str || !world_str) {
+        ctx->topology_invalid = true;
+        return -1;
+    }
+    const int rank = atoi(rank_str);
+    const int world_size = atoi(world_str);
+    const char* cp_size_str = getenv("CP_SIZE");
+    if (!cp_size_str) cp_size_str = getenv("RUSTRAIN_CP_SIZE");
+    const char* pp_size_str = getenv("PP_SIZE");
+    if (!pp_size_str) pp_size_str = getenv("RUSTRAIN_PP_SIZE");
+    if ((cp_size_str && atoi(cp_size_str) != 1) ||
+        (pp_size_str && atoi(pp_size_str) != 1)) {
+        fprintf(stderr,
+            "[parallel_nccl] qwen36_init_nccl only supports singleton CP/PP; "
+            "use qwen36_init_parallel_nccl_v2\n");
+        ctx->topology_invalid = true;
+        return -1;
+    }
+    const char* tp_size_str = getenv("TP_SIZE");
+    if (!tp_size_str) tp_size_str = getenv("RUSTRAIN_TP_SIZE");
+    const int tp_size = tp_size_str ? atoi(tp_size_str) : 1;
+    const bool data_parallel = env_enabled("RUSTRAIN_DATA_PARALLEL");
+    const char* ep_size_str = getenv("EP_SIZE");
+    if (!ep_size_str) ep_size_str = getenv("RUSTRAIN_EP_SIZE");
+    const int ep_size = ep_size_str ? atoi(ep_size_str) : 1;
+    const char* dp_size_str = getenv("DP_SIZE");
+    if (!dp_size_str) dp_size_str = getenv("RUSTRAIN_DP_SIZE");
+    const int dp_size = dp_size_str
+        ? atoi(dp_size_str)
+        : (data_parallel && tp_size > 0 && ep_size > 0
+            ? world_size / (tp_size * ep_size) : 1);
+    const int tp_rank = tp_size > 0 ? rank % tp_size : 0;
+    const int ep_rank = tp_size > 0 && ep_size > 0
+        ? (rank / tp_size) % ep_size : 0;
+    const int dp_rank = tp_size > 0 && ep_size > 0
+        ? rank / (tp_size * ep_size) : 0;
+    return qwen36_init_parallel_nccl_v2(
+        ctx_ptr, rank, world_size,
+        tp_rank, tp_size, rank / std::max(tp_size, 1),
+        0, 1, 0,
+        ep_rank, ep_size, dp_rank * std::max(tp_size, 1) + tp_rank,
+        dp_rank, dp_size, ep_rank * std::max(tp_size, 1) + tp_rank,
+        0, 1, 0);
+}
 
-            // Temporarily set lora_batch_valid so prepare_lora_batch runs
-            // We need to select only adapters[start:end]
-            // HACK: move non-chunk adapters to a temp vector, run, then restore
-            std::vector<TrainingContext::LoRAAdapter> temp;
-            temp.swap(ctx->adapters);
-            ctx->adapters.assign(temp.begin() + start, temp.begin() + end);
+// Set NCCL communicator for Expert Parallel all-reduce (legacy, from Rust)
+__attribute__((visibility("default"))) void qwen36_set_nccl_comm(
+    void* ctx_ptr, void* comm_ptr, void* stream_ptr,
+    int32_t ep_rank, int32_t ep_world_size
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    ctx->data_parallel = env_enabled("RUSTRAIN_DATA_PARALLEL");
+    if (ctx->tp_world_size <= 0 || ctx->tp_world_size > 1 ||
+        ctx->cp_world_size != 1 || ctx->pp_world_size != 1) {
+        ctx->topology_invalid = true;
+        fprintf(stderr,
+            "[parallel_nccl] legacy setter only supports singleton TP/CP/PP: "
+            "TP_SIZE=%d CP_SIZE=%d PP_SIZE=%d WORLD_SIZE=%d DATA_PARALLEL=%d\n",
+            ctx->tp_world_size, ctx->cp_world_size, ctx->pp_world_size,
+            ep_world_size, ctx->data_parallel ? 1 : 0);
+        return;
+    }
+    if (ctx->data_parallel) {
+        ctx->dp_comm = reinterpret_cast<ncclComm_t>(comm_ptr);
+        ctx->dp_stream = reinterpret_cast<cudaStream_t>(stream_ptr);
+        ctx->dp_rank = ep_rank;
+        ctx->dp_world_size = ep_world_size;
+        ctx->nccl_comm = nullptr;
+        ctx->nccl_stream = nullptr;
+        ctx->ep_rank = 0;
+        ctx->ep_world_size = 1;
+        ctx->expert_parallel = false;
+    } else {
+        ctx->nccl_comm = reinterpret_cast<ncclComm_t>(comm_ptr);
+        ctx->nccl_stream = reinterpret_cast<cudaStream_t>(stream_ptr);
+        ctx->ep_rank = ep_rank;
+        ctx->ep_world_size = ep_world_size;
+        ctx->expert_parallel = ep_world_size > 1;
+        ctx->dp_comm = nullptr;
+        ctx->dp_stream = nullptr;
+        ctx->dp_rank = 0;
+        ctx->dp_world_size = 1;
+    }
+    ctx->topology_invalid = false;
+    int current_device = g_cuda_device;
+    cudaGetDevice(&current_device);
+    ctx->cuda_device = current_device;
+    // Only EP owns routed-output collectives. In pure replicated DP the world
+    // communicator belongs exclusively to LoRA gradient synchronization;
+    // exposing it to moe_forward would mix activations from unrelated samples.
+    void* layer_comm = ctx->expert_parallel ? comm_ptr : nullptr;
+    void* layer_stream = ctx->expert_parallel ? stream_ptr : nullptr;
+    for (auto& lc : ctx->layer_configs) {
+        lc.nccl_comm = layer_comm;
+        lc.nccl_stream = layer_stream;
+    }
+    for (auto& lc : ctx->mtp_layer_configs) {
+        lc.nccl_comm = layer_comm;
+        lc.nccl_stream = layer_stream;
+    }
+}
 
-            // Mark batched mode active
-            ctx->lora_batch_valid = true;  // triggers prepare_lora_batch in forward
+// Enable/disable gradient checkpointing
+__attribute__((visibility("default"))) void qwen36_set_checkpoint(void* ctx_ptr, int32_t enable, int64_t group_size) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    ctx->use_checkpoint = (enable != 0);
+    ctx->group_size = (group_size > 0) ? group_size : 4;
+    fprintf(stderr, "[q36_ctx] checkpoint: %s, group_size=%ld\n",
+        ctx->use_checkpoint ? "ON" : "OFF", (long)ctx->group_size);
+}
 
-            // Expand input_ids from [1, seq] to [N, seq]
-            auto ids_expanded = input_ids.repeat({n, 1});
-            auto mask_expanded = target_mask.repeat({n, 1});
+// Set attention mask for padding tokens
+__attribute__((visibility("default"), used))
+void qwen36_set_attention_mask(void* ctx_ptr, void* mask_ptr) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (mask_ptr) {
+        auto& mask = *reinterpret_cast<at::Tensor*>(mask_ptr);
+        validate_linear_attention_mask(ctx, mask);
+        ctx->attention_mask = mask;
+        elide_trivial_attention_mask(ctx);
+    } else {
+        ctx->attention_mask = at::Tensor();
+        ctx->attention_lengths = at::Tensor();
+    }
+}
 
-            // Run train_step (reuses existing forward + loss + backward + Adam)
-            // But we need to pass the expanded tensors
-            auto& input_ref = ids_expanded;
-            auto& mask_ref = mask_expanded;
+// Utility functions (kept for compatibility)
+__attribute__((visibility("default"))) void* qwen36_gemm(void* a_ptr, void* b_ptr, int transpose_b) {
+    auto& a = *reinterpret_cast<at::Tensor*>(a_ptr);
+    auto& b = *reinterpret_cast<at::Tensor*>(b_ptr);
+    if (transpose_b) return new at::Tensor(at::matmul(a, b.t()));
+    return new at::Tensor(at::matmul(a, b));
+}
 
-            // Forward — force checkpoint for multi-LoRA (needed for group_inputs)
-            bool use_fused = getenv("QWEN36_FUSED_LAYER");
-            ctx->use_checkpoint = true;  // force checkpoint for manual_group_backward
+__attribute__((visibility("default"))) void qwen36_free_tensor(void* tensor_ptr) {
+    if (tensor_ptr) delete reinterpret_cast<at::Tensor*>(tensor_ptr);
+}
 
-            auto t_fwd_start = std::chrono::steady_clock::now();
-            auto hidden = use_fused
-                ? forward_full_fused(ctx, input_ref)
-                : forward_full_checkpoint(ctx, input_ref);
-            auto t_fwd_end = std::chrono::steady_clock::now();
-            double fwd_ms = std::chrono::duration<double, std::milli>(t_fwd_end - t_fwd_start).count();
+// ── Multi-LoRA adapter management ──
 
-            // Batched CE: compute loss with autograd enabled.
-            double loss_val;
-            auto t_loss_start = std::chrono::steady_clock::now();
-            {
-                at::AutoGradMode grad_enable(true);
-                // Re-attach hidden to autograd graph
-                hidden = hidden.detach().set_requires_grad(true);
-                if (getenv("QWEN36_FUSED_CE")) {
-                    auto loss = compute_loss_fused(ctx, hidden, input_ref, mask_ref, ctx->vocab_size);
-                    loss_val = loss.item<double>();
-                } else {
-                    auto loss = compute_loss(ctx, hidden, input_ref, mask_ref, ctx->vocab_size);
-                    loss_val = loss.item<double>();
+static int64_t qwen36_add_lora_impl(
+    void* ctx_ptr,
+    int64_t rank, double alpha,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str,
+    double optimizer_lr, double optimizer_beta1,
+    double optimizer_beta2, double optimizer_eps
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "dynamic LoRA registration requires a context");
+        validate_dynamic_context_health(ctx);
+        TrainingContext::LoRAAdapter adapter{};
+        int64_t local_rank = 0;
+        std::string request_error;
+        try {
+            require_canonical_adapter_index(ctx);
+            TORCH_CHECK(!ctx->topology_invalid,
+                "native Qwen context rejected an incompatible distributed "
+                "topology");
+            TORCH_CHECK(!ctx->accumulation_active &&
+                    ctx->accumulated_token_weight == 0.0 &&
+                    !ctx->pp_window.active,
+                "cannot mutate the dynamic LoRA registry while a gradient or "
+                "pipeline window is pending; finalize or abort it first");
+            TORCH_CHECK(ctx->router_aux_loss_coef == 0.0,
+                "dynamic multi-LoRA requires router_aux_loss_coef=0 until "
+                "per-tenant routing statistics are implemented");
+            TORCH_CHECK(rank > 0, "LoRA rank must be positive");
+            TORCH_CHECK(std::isfinite(alpha) && alpha > 0.0,
+                "LoRA alpha must be finite and positive");
+            TORCH_CHECK(std::isfinite(optimizer_lr) && optimizer_lr >= 0.0,
+                "dynamic LoRA optimizer learning rate must be finite and "
+                "non-negative");
+            TORCH_CHECK(std::isfinite(optimizer_beta1) &&
+                    optimizer_beta1 >= 0.0 && optimizer_beta1 < 1.0,
+                "dynamic LoRA optimizer beta1 must be finite and in [0, 1)");
+            TORCH_CHECK(std::isfinite(optimizer_beta2) &&
+                    optimizer_beta2 >= 0.0 && optimizer_beta2 < 1.0,
+                "dynamic LoRA optimizer beta2 must be finite and in [0, 1)");
+            const float beta1_f = static_cast<float>(optimizer_beta1);
+            const float beta2_f = static_cast<float>(optimizer_beta2);
+            TORCH_CHECK(std::isfinite(beta1_f) && beta1_f >= 0.0f &&
+                    beta1_f < 1.0f && std::isfinite(beta2_f) &&
+                    beta2_f >= 0.0f && beta2_f < 1.0f,
+                "dynamic LoRA optimizer betas must be representable as finite "
+                "FP32 values below 1");
+            TORCH_CHECK(std::isfinite(optimizer_eps) && optimizer_eps >= 0.0,
+                "dynamic LoRA optimizer epsilon must be finite and non-negative");
+            TORCH_CHECK(num_target_layers >= 0 &&
+                    (num_target_layers == 0 || target_layers),
+                "dynamic LoRA target layer list is invalid");
+            TORCH_CHECK(ctx->next_adapter_id <
+                    std::numeric_limits<int64_t>::max(),
+                "dynamic LoRA adapter ID space is exhausted");
+            adapter.id = ctx->next_adapter_id + 1;
+            adapter.rank = rank;
+            adapter.optimizer_lr = optimizer_lr;
+            // The fused state update consumes FP32 scalar buffers. Store the
+            // representable values so bias correction and recurrence use the
+            // same beta semantics, including near-one configurations.
+            adapter.optimizer_beta1 = beta1_f;
+            adapter.optimizer_beta2 = beta2_f;
+            adapter.optimizer_eps = optimizer_eps;
+            adapter.alpha = alpha;
+            adapter.all_target_layers = num_target_layers == 0;
+            for (int64_t i = 0; i < num_target_layers; i++) {
+                const int64_t global_layer = target_layers[i];
+                TORCH_CHECK(global_layer >= 0 &&
+                        global_layer < ctx->global_num_layers,
+                    "dynamic LoRA target layer out of range: ", global_layer,
+                    " for model with ", ctx->global_num_layers, " layers");
+                adapter.global_target_layers.insert(global_layer);
+                if (global_layer >= ctx->global_layer_start &&
+                        global_layer <
+                            ctx->global_layer_start + ctx->num_layers) {
+                    adapter.target_layers.insert(
+                        global_layer - ctx->global_layer_start);
                 }
             }
-            auto t_loss_end = std::chrono::steady_clock::now();
-            double loss_ms = std::chrono::duration<double, std::milli>(t_loss_end - t_loss_start).count();
-
-            // Backward
-            auto t_bwd_start = std::chrono::steady_clock::now();
-            auto hidden_grad = hidden.grad();
-            hidden = hidden.detach();
-            if (hidden.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache();
-
-            if (use_fused && hidden_grad.defined()) {
-                hidden.backward(hidden_grad);
-            } else if (hidden_grad.defined() && !ctx->group_inputs.empty() && !getenv("QWEN36_SUBCKPT")) {
-                manual_group_backward(ctx, hidden_grad);
-            } else if (hidden_grad.defined()) {
-                hidden.backward(hidden.grad());
+            if (target_modules_str) {
+                std::string s(target_modules_str);
+                std::stringstream ss(s);
+                std::string item;
+                while (std::getline(ss, item, ','))
+                    adapter.target_modules.insert(item);
             }
-            // No cudaDeviceSynchronize — let GPU pipeline run asynchronously.
-            // The next chunk's CPU prep (LoRA batch, input expand) will overlap
-            // with the tail of this chunk's GPU backward.
-            auto t_bwd_end = std::chrono::steady_clock::now();
-            double bwd_ms = std::chrono::duration<double, std::milli>(t_bwd_end - t_bwd_start).count();
-
-            fprintf(stderr, "[train_multi] chunk %ld/%ld: n=%ld loss=%f  fwd=%.0fms loss=%.0fms bwd=%.0fms\n",
-                    (long)(chunk+1), (long)num_chunks, (long)n, loss_val, fwd_ms, loss_ms, bwd_ms);
-
-            // MTP (if enabled)
-            if (ctx->has_mtp && !getenv("QWEN36_DISABLE_MTP")) {
-                auto hidden_detached = hidden.detach().set_requires_grad(true);
-                auto mtp_hidden = mtp_forward(ctx, hidden_detached, input_ref);
-                auto mtp_loss = mtp_compute_loss(ctx, mtp_hidden, input_ref, mask_ref);
-                mtp_loss.backward();
-                if (hidden_detached.grad().defined()) {
-                    if (hidden.grad().defined()) {
-                        hidden.grad().add_(hidden_detached.grad());
-                    } else {
-                        hidden.mutable_grad() = hidden_detached.grad().clone();
+            // The activation-level batch path stacks A/B across adapters.
+            // Keep the batch rectangular and semantically aligned instead of
+            // waiting for an opaque stack failure during the first step.
+            if (!ctx->adapters.empty() &&
+                !ctx->restore_without_parameter_sync &&
+                !ctx->allow_heterogeneous_registration) {
+                const auto& reference = ctx->adapters.front();
+                TORCH_CHECK(rank == reference.rank,
+                    "dynamic LoRA adapters in one batch must use the same rank");
+                TORCH_CHECK(adapter.all_target_layers ==
+                        reference.all_target_layers &&
+                        adapter.global_target_layers ==
+                            reference.global_target_layers,
+                    "dynamic LoRA adapters in one batch must use identical "
+                    "target_layers");
+                TORCH_CHECK(adapter.target_modules == reference.target_modules,
+                    "dynamic LoRA adapters in one batch must use identical "
+                    "target_modules");
+            }
+            for (const auto& name : adapter.target_modules) {
+                TORCH_CHECK(
+                    name == "q_proj" || name == "k_proj" ||
+                    name == "v_proj" || name == "o_proj" ||
+                    name == "in_proj_qkv" || name == "in_proj_z" ||
+                    name == "in_proj_a" || name == "in_proj_b" ||
+                    name == "out_proj" || name == "gate_proj" ||
+                    name == "up_proj" || name == "down_proj" ||
+                    name == "shared_gate_proj" ||
+                    name == "shared_up_proj" ||
+                    name == "shared_down_proj" ||
+                    name == "experts_gate_up_proj" ||
+                    name == "experts_down_proj",
+                    "unsupported dynamic Qwen LoRA target module: ", name);
+                bool resolved = false;
+                for (const auto& layer_cfg : ctx->layer_configs) {
+                    auto table = lora_projection_table(layer_cfg);
+                    for (int64_t pair = 0; pair < table.count; ++pair) {
+                        if (name == table.entries[pair].name) {
+                            resolved = true;
+                            break;
+                        }
                     }
+                    if (resolved) break;
                 }
+                TORCH_CHECK(resolved ||
+                        ctx->global_num_layers != ctx->num_layers,
+                    "dynamic LoRA target module does not exist in this model: ",
+                    name);
             }
-
-            // Adam step
-            at::AutoGradMode guard(false);
-            ctx->step_count++;
-            ctx->lora_cache_valid = false;
-            ctx->lora_batch_valid = false;
-
-            double step_f = (double)ctx->step_count;
-            double bias_correction1 = 1.0 - std::pow(ctx->beta1, step_f);
-            double bias_correction2 = 1.0 - std::pow(ctx->beta2, step_f);
-            float lr_scaled = (float)(ctx->lr / bias_correction1);
-            float eps_scaled = (float)(ctx->eps / std::sqrt(bias_correction2));
-            float one_minus_b1 = (float)(1.0 - ctx->beta1);
-            float one_minus_b2 = (float)(1.0 - ctx->beta2);
-
-            std::vector<void*> h_params, h_grads;
-            std::vector<float*> h_m, h_v;
-            std::vector<int> h_sizes;
-
-            for (auto& adapter : ctx->adapters) {
-                for (auto& [layer_idx, pairs] : adapter.params) {
-                    auto& adam_states = adapter.adam_state[layer_idx];
-                    for (size_t i = 0; i < pairs.size(); i++) {
-                        auto& [a, b] = pairs[i];
-                        auto& [m_a, v_a, m_b, v_b] = adam_states[i];
-                        if (a.grad().defined() && a.scalar_type() == at::kBFloat16) {
-                            h_params.push_back(a.data_ptr());
-                            h_grads.push_back(a.grad().data_ptr());
-                            h_m.push_back((float*)m_a.data_ptr());
-                            h_v.push_back((float*)v_a.data_ptr());
-                            h_sizes.push_back((int)a.numel());
-                        }
-                        if (b.grad().defined() && b.scalar_type() == at::kBFloat16) {
-                            h_params.push_back(b.data_ptr());
-                            h_grads.push_back(b.grad().data_ptr());
-                            h_m.push_back((float*)m_b.data_ptr());
-                            h_v.push_back((float*)v_b.data_ptr());
-                            h_sizes.push_back((int)b.numel());
+            local_rank = local_lora_rank_for_active_targets(
+                ctx, rank, adapter.target_layers, adapter.target_modules,
+                adapter.all_target_layers,
+                /*empty_modules_mean_attention_only=*/true, "dynamic");
+        } catch (const std::exception& e) {
+            request_error = e.what();
+        }
+        const bool local_request_valid = request_error.empty();
+        const bool request_matches = adapter_registration_phase_matches(
+            ctx, adapter, local_request_valid, /*phase=*/0);
+        TORCH_CHECK(request_matches && local_request_valid,
+            "dynamic LoRA registration request is invalid or differs across "
+            "distributed ranks", request_error.empty() ? "" : ": ",
+            request_error);
+
+        std::string preparation_error;
+        try {
+            for (int64_t i = 0; i < ctx->num_layers; i++) {
+                if (!adapter.all_target_layers &&
+                    adapter.target_layers.find(i) ==
+                        adapter.target_layers.end())
+                    continue;
+                int64_t w_offset = 0;
+                for (int64_t j = 0; j < i; j++)
+                    w_offset += weight_count_for_layer(ctx->layer_configs[j]);
+                auto projection_table =
+                    lora_projection_table(ctx->layer_configs[i]);
+                int64_t num_pairs = projection_table.count;
+                std::vector<std::pair<at::Tensor, at::Tensor>> pairs;
+                std::vector<std::array<at::Tensor, 4>> adam_states;
+                std::vector<std::array<at::Tensor, 6>> adam_shadows;
+                std::vector<int64_t> adam_shadow_slots;
+                std::vector<std::array<at::Tensor, 2>> grad_accumulators;
+                const bool lazy_optimizer_state = env_enabled(
+                    "QWEN36_LAZY_DYNAMIC_OPTIMIZER_STATE", true);
+                const bool pooled_optimizer_shadow = env_enabled(
+                    "QWEN36_DYNAMIC_ADAM_SHADOW_POOL", true);
+                for (int64_t k = 0; k < num_pairs; k++) {
+                    const auto& projection = projection_table.entries[k];
+                    auto* base =
+                        ctx->weight_ptrs[w_offset + projection.weight_index];
+                    // Preserve the historical empty-target default (attention
+                    // projections only). Explicit lists may select any 2D
+                    // dense/shared MLP projection.
+                    bool active = adapter.target_modules.empty()
+                        ? !projection.grouped_expert &&
+                            projection.segment == LoraSegment::Attention
+                        : adapter.target_modules.find(projection.name) !=
+                            adapter.target_modules.end();
+                    auto opts = at::TensorOptions().dtype(ctx->compute_type)
+                        .device(base->device());
+                    at::Tensor a, b;
+                    if (active) {
+                        if (projection.grouped_expert) {
+                            TORCH_CHECK(base->dim() == 3,
+                                "dynamic routed-expert LoRA projection must "
+                                "be rank 3: ", projection.name);
+                            int64_t experts = base->size(0);
+                            int64_t out_f = base->size(1);
+                            int64_t in_f = base->size(2);
+                            const auto layout = lora_tp_layout(ctx, i, k);
+                            if (layout == LoraTpLayout::ColumnParallel ||
+                                layout == LoraTpLayout::RowParallel) {
+                                a = at::randn(
+                                    {experts, rank, in_f}, opts) * 0.01;
+                                b = at::zeros(
+                                    {experts, out_f, rank}, opts);
+                            } else {
+                                a = initialize_lora_a(
+                                    ctx, opts, experts, rank, in_f);
+                                b = at::zeros(
+                                    {experts, out_f, local_rank}, opts);
+                            }
+                        } else {
+                            TORCH_CHECK(base->dim() == 2,
+                                "dynamic LoRA projection must be a matrix: ",
+                                projection.name);
+                            int64_t out_f = base->size(0);
+                            int64_t in_f = base->size(1);
+                            const auto layout = lora_tp_layout(ctx, i, k);
+                            if (layout == LoraTpLayout::ColumnParallel ||
+                                layout == LoraTpLayout::RowParallel) {
+                                a = at::randn({rank, in_f}, opts) * 0.01;
+                                b = at::zeros({out_f, rank}, opts);
+                            } else {
+                                a = initialize_lora_a(
+                                    ctx, opts, 0, rank, in_f);
+                                b = at::zeros({out_f, local_rank}, opts);
+                            }
                         }
+                    } else {
+                        a = at::zeros({}, opts);
+                        b = at::zeros({}, opts);
                     }
+                    a.set_requires_grad(active);
+                    b.set_requires_grad(active);
+                    auto opts_f32 = at::TensorOptions().dtype(at::kFloat)
+                        .device(base->device());
+                    adam_states.push_back(lazy_optimizer_state
+                        ? std::array<at::Tensor, 4>{}
+                        : std::array<at::Tensor, 4>{
+                            at::zeros(a.sizes(), opts_f32),
+                            at::zeros(a.sizes(), opts_f32),
+                            at::zeros(b.sizes(), opts_f32),
+                            at::zeros(b.sizes(), opts_f32)});
+                    adam_shadows.push_back(
+                        active && !lazy_optimizer_state &&
+                                !pooled_optimizer_shadow
+                        ? std::array<at::Tensor, 6>{
+                            at::empty_like(a).set_requires_grad(true),
+                            at::empty(a.sizes(), opts_f32),
+                            at::empty(a.sizes(), opts_f32),
+                            at::empty_like(b).set_requires_grad(true),
+                            at::empty(b.sizes(), opts_f32),
+                            at::empty(b.sizes(), opts_f32)}
+                        : std::array<at::Tensor, 6>{});
+                    adam_shadow_slots.push_back(-1);
+                    grad_accumulators.push_back(
+                        std::array<at::Tensor, 2>{
+                            at::Tensor(), at::Tensor()});
+                    pairs.emplace_back(std::move(a), std::move(b));
                 }
+                adapter.params[i] = std::move(pairs);
+                adapter.adam_state[i] = std::move(adam_states);
+                adapter.adam_shadow[i] = std::move(adam_shadows);
+                adapter.adam_shadow_slots[i] =
+                    std::move(adam_shadow_slots);
+                adapter.grad_accum[i] = std::move(grad_accumulators);
             }
-
-            if (!h_params.empty()) {
-                int n_params = (int)h_params.size();
-                auto opts_cpu_long = at::TensorOptions().dtype(at::kLong).device(at::kCPU);
-                auto opts_cpu_int  = at::TensorOptions().dtype(at::kInt).device(at::kCPU);
-                auto params_cpu = at::from_blob(h_params.data(), {n_params}, opts_cpu_long);
-                auto grads_cpu  = at::from_blob(h_grads.data(),  {n_params}, opts_cpu_long);
-                auto m_cpu      = at::from_blob(h_m.data(),      {n_params}, opts_cpu_long);
-                auto v_cpu      = at::from_blob(h_v.data(),      {n_params}, opts_cpu_long);
-                auto sizes_cpu  = at::from_blob(h_sizes.data(),  {n_params}, opts_cpu_int);
-                ctx->adam_dev_bufs.ensure(n_params, ctx->adapters[0].params.begin()->second[0].first);
-                ctx->adam_dev_bufs.params_buf.copy_(params_cpu);
-                ctx->adam_dev_bufs.grads_buf.copy_(grads_cpu);
-                ctx->adam_dev_bufs.m_buf.copy_(m_cpu);
-                ctx->adam_dev_bufs.v_buf.copy_(v_cpu);
-                ctx->adam_dev_bufs.sizes_buf.copy_(sizes_cpu);
-
-                auto stream = c10::cuda::getCurrentCUDAStream().stream();
-                launch_fused_adam_multi(
-                    (void**)ctx->adam_dev_bufs.params_buf.data_ptr(),
-                    (void**)ctx->adam_dev_bufs.grads_buf.data_ptr(),
-                    (float**)ctx->adam_dev_bufs.m_buf.data_ptr(),
-                    (float**)ctx->adam_dev_bufs.v_buf.data_ptr(),
-                    (int*)ctx->adam_dev_bufs.sizes_buf.data_ptr(),
-                    n_params,
-                    (float)ctx->beta1, (float)ctx->beta2,
-                    lr_scaled, eps_scaled,
-                    one_minus_b1, one_minus_b2,
-                    (void*)stream
-                );
-            }
-
-            // Restore adapters
-            ctx->adapters.swap(temp);
-            // Re-merge: put updated chunk adapters back
-            for (int64_t i = 0; i < n; i++) {
-                temp[start + i] = std::move(ctx->adapters[i]);
-            }
-            ctx->adapters.swap(temp);
-
-            total_loss += loss_val;
-            // emptyCache at chunk boundary — does NOT sync GPU.
-            c10::cuda::CUDACachingAllocator::emptyCache();
-
-            fprintf(stderr, "[train_multi] chunk %ld/%ld: n=%ld loss=%.6f\n",
-                    (long)(chunk + 1), (long)num_chunks, (long)n, loss_val);
-        }
-
-        return total_loss / num_chunks;
+            bind_adapter_lora_gradient_slab(ctx, adapter);
+            ctx->adapters.reserve(ctx->adapters.size() + 1);
+            ctx->adapter_id_index.reserve(ctx->adapter_id_index.size() + 1);
+        } catch (const std::exception& e) {
+            preparation_error = e.what();
+        }
+        const bool local_prepared = preparation_error.empty();
+        const bool preparation_matches = adapter_registration_phase_matches(
+            ctx, adapter, local_prepared, /*phase=*/1);
+        TORCH_CHECK(preparation_matches && local_prepared,
+            "dynamic LoRA registration preparation failed or differs across "
+            "distributed ranks", preparation_error.empty() ? "" : ": ",
+            preparation_error);
+        std::string synchronization_error;
+        try {
+            if (!ctx->restore_without_parameter_sync)
+                synchronize_adapter_replicated_lora_parameters(ctx, adapter);
+            TORCH_CHECK(!std::getenv(
+                    "QWEN36_TEST_FAIL_ADAPTER_REGISTRATION_AFTER_SYNC"),
+                "injected dynamic LoRA registration failure after "
+                "parameter synchronization");
+        } catch (const std::exception& e) {
+            synchronization_error = e.what();
+        }
+        const bool local_synchronized = synchronization_error.empty();
+        const bool synchronization_matches = adapter_registration_phase_matches(
+            ctx, adapter, local_synchronized, /*phase=*/2);
+        TORCH_CHECK(synchronization_matches && local_synchronized,
+            "dynamic LoRA parameter synchronization failed on at least one "
+            "distributed rank",
+            synchronization_error.empty() ? "" : ": ",
+            synchronization_error);
+        int64_t id = adapter.id;
+        const size_t canonical_index = ctx->adapters.size();
+        ctx->adapters.push_back(std::move(adapter));
+        ctx->adapter_id_index.emplace_back(id, canonical_index);
+        ctx->next_adapter_id = id;
+        ctx->lora_cache_valid = false;
+        ctx->lora_batch_valid = false;
+        fprintf(stderr,
+            "[q36_lora] added adapter %ld: rank=%ld alpha=%.1f "
+            "lr=%.8g beta1=%.8g beta2=%.8g eps=%.8g\n",
+            (long)id, (long)rank, alpha, optimizer_lr,
+            optimizer_beta1, optimizer_beta2, optimizer_eps);
+        return id;
     } catch (const std::exception& e) {
-        fprintf(stderr, "[train_multi] FAILED: %s\n", e.what());
-        return -1.0;
-    } catch (...) {
-        fprintf(stderr, "[train_multi] FAILED: unknown exception\n");
-        return -1.0;
+        fprintf(stderr, "[q36] add_lora FAILED: %s\n", e.what());
+        return -1;
     }
 }
 
-// Set NCCL communicator for Expert Parallel all-reduce
-// Creates NCCL communicator directly in C++ using env vars RANK/WORLD_SIZE.
-// Rank 0 generates unique ID and writes to /tmp/rustrain-nccl/nccl-id.bin
-// Other ranks read it. All ranks call ncclCommInitRank.
-// Returns 0 on success, -1 on failure.
-// Process-level NCCL singleton — created once, reused across sessions.
-static ncclComm_t g_nccl_comm = nullptr;
-static cudaStream_t g_nccl_stream = nullptr;
-static bool g_nccl_initialized = false;
+struct ScopedBoolOverride {
+    bool& target;
+    bool previous;
 
-// Set CUDA device — called from Rust worker before any GPU operation.
-// Ensures PyTorch initializes CUDA context on the correct device.
-__attribute__((visibility("default"))) void qwen36_set_cuda_device(int32_t device) {
-    // Use PyTorch's device API — this updates both cudaSetDevice AND
-    // PyTorch's internal device tracking (c10::cuda::current_device).
-    // Must be called before any GPU operation in exec'd worker processes.
-    c10::cuda::set_device(device);
-    cudaSetDevice(device);
-    // Force PyTorch to create CUDA context on this device
-    auto opts = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, device);
-    auto dummy = at::empty({1}, opts);
-    dummy.sizes();  // touch to ensure materialization
+    ScopedBoolOverride(bool& target_value, bool value)
+        : target(target_value), previous(target_value) {
+        target = value;
+    }
+
+    ~ScopedBoolOverride() { target = previous; }
+};
+
+__attribute__((visibility("default")))
+int64_t qwen36_add_lora(
+    void* ctx_ptr,
+    int64_t rank, double alpha,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    return qwen36_add_lora_impl(
+        ctx_ptr, rank, alpha, target_layers, num_target_layers,
+        target_modules_str,
+        ctx ? ctx->lr : 0.0,
+        ctx ? ctx->beta1 : 0.0,
+        ctx ? ctx->beta2 : 0.0,
+        ctx ? ctx->eps : 0.0);
 }
 
-__attribute__((visibility("default"))) int32_t qwen36_init_nccl(
-    void* ctx_ptr
+__attribute__((visibility("default")))
+int64_t qwen36_add_lora_v2(
+    void* ctx_ptr,
+    int64_t rank, double alpha,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str
 ) {
     auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx) return -1;
+    ScopedBoolOverride guard(ctx->allow_heterogeneous_registration, true);
+    return qwen36_add_lora(
+        ctx_ptr, rank, alpha, target_layers, num_target_layers,
+        target_modules_str);
+}
 
-    // If already initialized, just set the pointer on this context
-    if (g_nccl_initialized) {
-        ctx->nccl_comm = g_nccl_comm;
-        ctx->nccl_stream = g_nccl_stream;
-        // CRITICAL: also set ep_rank/ep_world_size — needed for new TrainingContext
-        // created by subsequent CreateSession commands. The fast path previously
-        // skipped this, leaving ep_rank=0 → cudaSetDevice(0) on all ranks → crash.
-        const char* rank_str2 = getenv("RANK");
-        const char* world_str2 = getenv("WORLD_SIZE");
-        if (rank_str2) ctx->ep_rank = atoi(rank_str2);
-        if (world_str2) ctx->ep_world_size = atoi(world_str2);
-        for (auto& lc : ctx->layer_configs) { lc.nccl_comm = (void*)g_nccl_comm; lc.nccl_stream = (void*)g_nccl_stream; }
-        for (auto& lc : ctx->mtp_layer_configs) { lc.nccl_comm = (void*)g_nccl_comm; lc.nccl_stream = (void*)g_nccl_stream; }
+// Additive ABI28 extension. Older registration entry points preserve their
+// context-wide learning-rate behavior; callers that resolve this optional
+// symbol can select one learning rate per dynamic tenant.
+__attribute__((visibility("default")))
+int64_t qwen36_add_lora_with_optimizer(
+    void* ctx_ptr,
+    int64_t rank, double alpha,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str,
+    double optimizer_lr
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx) return -1;
+    ScopedBoolOverride guard(ctx->allow_heterogeneous_registration, true);
+    return qwen36_add_lora_impl(
+        ctx_ptr, rank, alpha, target_layers, num_target_layers,
+        target_modules_str, optimizer_lr,
+        ctx->beta1, ctx->beta2, ctx->eps);
+}
+
+// Additive ABI28 extension for complete per-tenant Adam isolation. The
+// transactional update still executes all selected tenants in one fused
+// kernel launch through per-tensor optimizer scalar buffers.
+__attribute__((visibility("default")))
+int64_t qwen36_add_lora_with_optimizer_v2(
+    void* ctx_ptr,
+    int64_t rank, double alpha,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str,
+    double optimizer_lr, double optimizer_beta1,
+    double optimizer_beta2, double optimizer_eps
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx) return -1;
+    ScopedBoolOverride guard(ctx->allow_heterogeneous_registration, true);
+    return qwen36_add_lora_impl(
+        ctx_ptr, rank, alpha, target_layers, num_target_layers,
+        target_modules_str, optimizer_lr, optimizer_beta1,
+        optimizer_beta2, optimizer_eps);
+}
+
+__attribute__((visibility("default")))
+int64_t qwen36_add_lora_for_restore(
+    void* ctx_ptr,
+    int64_t rank, double alpha,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx) return -1;
+    ScopedBoolOverride guard(ctx->restore_without_parameter_sync, true);
+    return qwen36_add_lora(
+        ctx_ptr, rank, alpha, target_layers, num_target_layers,
+        target_modules_str);
+}
+
+// Checkpoint restore variant for tenants with an adapter-specific learning
+// rate. Both heterogeneous registration and temporary-parameter sync are
+// scoped off for this allocation; the caller hydrates tensors immediately.
+__attribute__((visibility("default")))
+int64_t qwen36_add_lora_for_restore_with_optimizer(
+    void* ctx_ptr,
+    int64_t rank, double alpha,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str,
+    double optimizer_lr
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx) return -1;
+    ScopedBoolOverride restore_guard(ctx->restore_without_parameter_sync, true);
+    ScopedBoolOverride heterogeneous_guard(ctx->allow_heterogeneous_registration, true);
+    return qwen36_add_lora_impl(
+        ctx_ptr, rank, alpha, target_layers, num_target_layers,
+        target_modules_str, optimizer_lr,
+        ctx->beta1, ctx->beta2, ctx->eps);
+}
+
+__attribute__((visibility("default")))
+int64_t qwen36_add_lora_for_restore_with_optimizer_v2(
+    void* ctx_ptr,
+    int64_t rank, double alpha,
+    const int64_t* target_layers, int64_t num_target_layers,
+    const char* target_modules_str,
+    double optimizer_lr, double optimizer_beta1,
+    double optimizer_beta2, double optimizer_eps
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx) return -1;
+    ScopedBoolOverride restore_guard(ctx->restore_without_parameter_sync, true);
+    ScopedBoolOverride heterogeneous_guard(ctx->allow_heterogeneous_registration, true);
+    return qwen36_add_lora_impl(
+        ctx_ptr, rank, alpha, target_layers, num_target_layers,
+        target_modules_str, optimizer_lr, optimizer_beta1,
+        optimizer_beta2, optimizer_eps);
+}
+
+// Restore a dynamic adapter's externally visible ID during checkpoint load.
+// IDs are positive and unique; the monotonic allocator is advanced so future
+// additions cannot collide with a restored tenant.
+__attribute__((visibility("default")))
+int32_t qwen36_set_adapter_id(void* ctx_ptr, int64_t current_id, int64_t requested_id) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx && current_id > 0 && requested_id > 0,
+            "dynamic adapter IDs must be positive");
+        require_clear_accumulation_for_registry_mutation(ctx);
+        const bool index_available = ctx->adapter_id_index_valid &&
+            ctx->adapter_id_index.size() == ctx->adapters.size();
+        size_t current_index = 0;
+        size_t requested_index = 0;
+        const bool current_found = find_canonical_adapter_index(
+            ctx, current_id, current_index);
+        const bool requested_found = find_canonical_adapter_index(
+            ctx, requested_id, requested_index);
+        const bool requested_available =
+            !requested_found || requested_id == current_id;
+        TrainingContext::LoRAAdapter request{};
+        request.id = current_id;
+        request.rank = requested_id;
+        request.alpha = 1.0;
+        const bool local_valid =
+            index_available && current_found && requested_available;
+        TORCH_CHECK(adapter_registration_phase_matches(
+                ctx, request, local_valid, /*phase=*/3) && local_valid,
+            current_found ? "dynamic adapter ID already exists: " :
+                "dynamic adapter not found: ",
+            current_found ? requested_id : current_id);
+        auto index_it = std::lower_bound(
+            ctx->adapter_id_index.begin(), ctx->adapter_id_index.end(),
+            current_id,
+            [](const auto& entry, int64_t id) { return entry.first < id; });
+        TORCH_CHECK(index_it != ctx->adapter_id_index.end() &&
+                index_it->first == current_id &&
+                index_it->second == current_index,
+            "validated dynamic adapter disappeared: ", current_id);
+        ctx->adapters[current_index].id = requested_id;
+        index_it->first = requested_id;
+        std::sort(
+            ctx->adapter_id_index.begin(), ctx->adapter_id_index.end(),
+            [](const auto& left, const auto& right) {
+                return left.first < right.first;
+            });
+        ctx->next_adapter_id = std::max(ctx->next_adapter_id, requested_id);
+        ctx->lora_cache_valid = false;
+        ctx->lora_batch_valid = false;
         return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] set_adapter_id FAILED: %s\n", e.what());
+        return -1;
     }
+}
 
-    const char* rank_str = getenv("RANK");
-    const char* world_str = getenv("WORLD_SIZE");
-    if (!rank_str || !world_str) return -1;
-    int rank = atoi(rank_str);
-    int world_size = atoi(world_str);
-    if (world_size <= 1) return 0;  // no EP needed
+__attribute__((visibility("default")))
+int32_t qwen36_remove_lora(void* ctx_ptr, int64_t adapter_id) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx && adapter_id > 0,
+            "dynamic LoRA removal requires a context and positive adapter ID");
+        require_clear_accumulation_for_registry_mutation(ctx);
+        const bool index_available = ctx->adapter_id_index_valid &&
+            ctx->adapter_id_index.size() == ctx->adapters.size();
+        size_t canonical_index = 0;
+        const bool local_found = find_canonical_adapter_index(
+            ctx, adapter_id, canonical_index);
+        TrainingContext::LoRAAdapter request{};
+        request.id = adapter_id;
+        request.rank = index_available ? -1 : -2;
+        request.optimizer_step = local_found ? 1 : 0;
+        request.alpha = 1.0;
+        TORCH_CHECK(adapter_registration_phase_matches(
+                ctx, request, index_available, /*phase=*/4),
+            "dynamic LoRA removal request or registry differs across ranks");
+        if (!local_found) return 0;
+        const auto index_it = std::lower_bound(
+            ctx->adapter_id_index.begin(), ctx->adapter_id_index.end(),
+            adapter_id,
+            [](const auto& entry, int64_t id) { return entry.first < id; });
+        TORCH_CHECK(index_it != ctx->adapter_id_index.end() &&
+                index_it->first == adapter_id &&
+                index_it->second == canonical_index,
+            "validated dynamic adapter disappeared: ", adapter_id);
+        ctx->adapters.erase(ctx->adapters.begin() + canonical_index);
+        ctx->adapter_id_index.erase(index_it);
+        for (auto& entry : ctx->adapter_id_index) {
+            if (entry.second > canonical_index) --entry.second;
+        }
+        ctx->lora_cache_valid = false;
+        ctx->lora_batch_valid = false;
+        return 1;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] remove_lora FAILED: %s\n", e.what());
+        return -1;
+    }
+}
 
-    // Set CUDA device and initialize PyTorch CUDA context on this device.
-    // NCCL communicator binds to the current CUDA context. If PyTorch hasn't
-    // initialized CUDA on this device yet, NCCL gets a wrong context → "invalid argument".
-    // Creating a dummy tensor forces PyTorch to initialize its CUDA context.
-    const char* local_rank_str = getenv("LOCAL_RANK");
-    int local_rank = local_rank_str ? atoi(local_rank_str) : rank;
-    cudaSetDevice(local_rank);
-    {
-        // Force PyTorch CUDA context initialization on this device
-        auto opts = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, local_rank);
-        auto dummy = at::empty({1}, opts);
-        dummy.sizes();  // touch to ensure materialization
+__attribute__((visibility("default")))
+int64_t qwen36_list_lora(void* ctx_ptr, int64_t* out_ids, int64_t max_count) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx || max_count < 0) return -1;
+    int64_t count = (int64_t)ctx->adapters.size();
+    // A zero-capacity call is a count query. This keeps the existing ABI while
+    // allowing callers to enumerate registries larger than the old 64-entry
+    // convenience buffer.
+    if (max_count == 0) return count;
+    if (!out_ids) return -1;
+    if (count > max_count) count = max_count;
+    for (int64_t i = 0; i < count; i++)
+        out_ids[i] = ctx->adapters[i].id;
+    return count;
+}
+
+__attribute__((visibility("default")))
+void* qwen36_get_adapter_lora_tensor(
+    void* ctx_ptr, int64_t adapter_id, int64_t global_layer,
+    const char* module_name, int32_t is_b
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    int64_t layer_idx = -1;
+    if (!module_name ||
+            !global_to_local_layer(ctx, global_layer, layer_idx)) return nullptr;
+    const int64_t pair_idx = lora_pair_index(
+        ctx->layer_configs[layer_idx], module_name);
+    if (pair_idx < 0) return nullptr;
+    for (auto& adapter : ctx->adapters) {
+        if (adapter.id != adapter_id) continue;
+        auto it = adapter.params.find(layer_idx);
+        if (it == adapter.params.end() ||
+            pair_idx >= static_cast<int64_t>(it->second.size()))
+            return nullptr;
+        auto& pair = it->second[pair_idx];
+        auto& tensor = is_b ? pair.second : pair.first;
+        return tensor.requires_grad() ? &tensor : nullptr;
     }
+    return nullptr;
+}
 
-    // Exchange unique ID via file
-    // Rank 0 generates ID, writes to file, then writes "ready" sentinel.
-    // Other ranks wait for "ready" file, then read ID.
-    // This ensures rank 0's write is visible before others read.
-    const char* id_path = "/tmp/rustrain-nccl/nccl-id.bin";
-    const char* ready_path = "/tmp/rustrain-nccl/nccl-ready.txt";
-    ncclUniqueId unique_id;
-    if (rank == 0) {
-        mkdir("/tmp/rustrain-nccl", 0777);
-        // Clean up old files first
-        remove("/tmp/rustrain-nccl/nccl-ready.txt");
-        ncclGetUniqueId(&unique_id);
-        FILE* f = fopen(id_path, "wb");
-        fwrite(&unique_id, sizeof(unique_id), 1, f);
-        fclose(f);
-        // Write ready sentinel AFTER id file
-        FILE* rf = fopen(ready_path, "w");
-        fprintf(rf, "ready\n");
-        fclose(rf);
-    } else {
-        // Wait for ready sentinel
-        for (int i = 0; i < 600; i++) {
-            FILE* rf = fopen(ready_path, "r");
-            if (rf) { fclose(rf); break; }
-            usleep(10000);  // 10ms
-        }
-        // Now read ID file
-        FILE* f = fopen(id_path, "rb");
-        if (!f || fread(&unique_id, sizeof(unique_id), 1, f) != 1) {
-            fprintf(stderr, "[ep_nccl] rank %d: failed to read ID file\n", rank);
-            if (f) fclose(f);
-            return -1;
-        }
-        fclose(f);
+__attribute__((visibility("default")))
+int32_t qwen36_set_adapter_lora_tensor(
+    void* ctx_ptr, int64_t adapter_id, int64_t global_layer,
+    const char* module_name, int32_t is_b, void* tensor_ptr
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx && module_name && tensor_ptr,
+            "invalid dynamic LoRA tensor setter arguments");
+        auto* target = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_lora_tensor(
+                ctx, adapter_id, global_layer, module_name, is_b));
+        TORCH_CHECK(target, "dynamic LoRA target not found: adapter=", adapter_id,
+            " global_layer=", global_layer, " module=", module_name);
+        auto& source = *reinterpret_cast<at::Tensor*>(tensor_ptr);
+        TORCH_CHECK(source.sizes() == target->sizes(),
+            "dynamic LoRA tensor shape mismatch: expected ", target->sizes(),
+            " got ", source.sizes());
+        at::NoGradGuard guard;
+        target->copy_(source.to(target->device()).to(target->scalar_type()));
+        ctx->lora_cache_valid = false;
+        ctx->lora_batch_valid = false;
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] set_adapter_lora_tensor FAILED: %s\n", e.what());
+        return -1;
     }
+}
 
-    // Barrier: ensure all ranks reach ncclCommInitRank simultaneously.
-    // Without this, rank 0 (fast load_model) reaches ncclCommInitRank before
-    // rank 3 (slow load_model) → NCCL timeout.
-    {
-        const char* barrier_dir = "/tmp/rustrain-nccl";
-        char bpath[256];
-        snprintf(bpath, sizeof(bpath), "%s/barrier_%d", barrier_dir, rank);
-        FILE* bf = fopen(bpath, "w"); fprintf(bf, "1\n"); fclose(bf);
-        // Wait for all ranks
-        for (int i = 0; i < world_size; i++) {
-            char p[256];
-            snprintf(p, sizeof(p), "%s/barrier_%d", barrier_dir, i);
-            for (int w = 0; w < 6000; w++) {  // 60s timeout
-                FILE* f2 = fopen(p, "r");
-                if (f2) { fclose(f2); break; }
-                usleep(10000);
+// Access one dynamic adapter's Adam state. The per-layer state is stored as
+// {m_a, v_a, m_b, v_b}; `is_b` selects A/B and `is_v` selects m/v.
+__attribute__((visibility("default")))
+void* qwen36_get_adapter_optimizer_tensor(
+    void* ctx_ptr, int64_t adapter_id, int64_t global_layer,
+    const char* module_name, int32_t is_b, int32_t is_v
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        int64_t layer_idx = -1;
+        if (!module_name ||
+                !global_to_local_layer(ctx, global_layer, layer_idx))
+            return nullptr;
+        const int64_t pair_idx = lora_pair_index(
+            ctx->layer_configs[layer_idx], module_name);
+        if (pair_idx < 0) return nullptr;
+        for (auto& adapter : ctx->adapters) {
+            if (adapter.id != adapter_id) continue;
+            if (env_enabled(
+                    "QWEN36_LAZY_DYNAMIC_OPTIMIZER_STATE", true)) {
+                materialize_dynamic_adam_state(adapter);
             }
+            auto state_it = adapter.adam_state.find(layer_idx);
+            if (state_it == adapter.adam_state.end() ||
+                pair_idx >= static_cast<int64_t>(state_it->second.size()))
+                return nullptr;
+            // array order: m_a, v_a, m_b, v_b
+            const int index = (is_b ? 2 : 0) + (is_v ? 1 : 0);
+            auto& tensor = state_it->second[pair_idx][index];
+            return tensor.defined() ? &tensor : nullptr;
         }
+        return nullptr;
+    } catch (const std::exception& e) {
+        fprintf(stderr,
+            "[q36] get_adapter_optimizer_tensor FAILED: %s\n", e.what());
+        return nullptr;
+    }
+}
+
+// Export a checkpoint-owned CPU snapshot without materializing lazy device
+// state. Status 0 returns a heap-allocated at::Tensor wrapper through `out`;
+// status 1 means the adapter has never completed an optimizer step; -1 is an
+// invalid or inconsistent request. The caller releases a successful result
+// with qwen36_free_tensor.
+__attribute__((visibility("default")))
+int32_t qwen36_export_adapter_optimizer_tensor_cpu_v1(
+    void* ctx_ptr, int64_t adapter_id, int64_t global_layer,
+    const char* module_name, int32_t is_b, int32_t is_v, void** out
+) {
+    if (out) *out = nullptr;
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx && out && module_name && adapter_id > 0,
+            "invalid dynamic Adam checkpoint export arguments");
+        TORCH_CHECK((is_b == 0 || is_b == 1) &&
+                (is_v == 0 || is_v == 1),
+            "dynamic Adam checkpoint selectors must be boolean");
+        TORCH_CHECK(!ctx->dynamic_adam_transaction_active,
+            "cannot export dynamic Adam during an active transaction");
+        int64_t layer_idx = -1;
+        TORCH_CHECK(global_to_local_layer(ctx, global_layer, layer_idx),
+            "dynamic Adam checkpoint layer is not owned by this stage: ",
+            global_layer);
+        const int64_t pair_idx = lora_pair_index(
+            ctx->layer_configs[layer_idx], module_name);
+        TORCH_CHECK(pair_idx >= 0,
+            "dynamic Adam checkpoint module is not present: ", module_name);
+        size_t adapter_index = 0;
+        TORCH_CHECK(find_canonical_adapter_index(
+                ctx, adapter_id, adapter_index),
+            "dynamic Adam checkpoint adapter not found: ", adapter_id);
+        auto& adapter = ctx->adapters[adapter_index];
+        if (adapter.optimizer_step == 0) return 1;
+        snapshot_dynamic_adam_state_to_host(ctx, adapter);
+        auto state_it = adapter.adam_host_state.find(layer_idx);
+        TORCH_CHECK(state_it != adapter.adam_host_state.end() &&
+                pair_idx < static_cast<int64_t>(state_it->second.size()),
+            "dynamic Adam host snapshot layout mismatch for adapter ",
+            adapter_id, " layer ", global_layer);
+        const int index = (is_b ? 2 : 0) + (is_v ? 1 : 0);
+        const auto& tensor = state_it->second[pair_idx][index];
+        TORCH_CHECK(tensor.defined() && tensor.device().is_cpu() &&
+                tensor.is_contiguous() &&
+                tensor.scalar_type() == at::kFloat,
+            "dynamic Adam host snapshot tensor is unavailable for adapter ",
+            adapter_id, " layer ", global_layer, " module ", module_name);
+        *out = new at::Tensor(tensor);
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr,
+            "[q36] export_adapter_optimizer_tensor_cpu FAILED: %s\n",
+            e.what());
+        return -1;
     }
+}
 
-    // Initialize communicator — only once per process
-    ncclComm_t comm;
-    ncclResult_t err = ncclCommInitRank(&comm, world_size, unique_id, rank);
-    if (err != ncclSuccess) {
-        fprintf(stderr, "[ep_nccl] ncclCommInitRank failed: %d (%s)\n", err, ncclGetErrorString(err));
+// Atomically restore one tenant's complete stage-local Adam state into host
+// memory. Each slot contributes [m_a, v_a, m_b, v_b]. Device state remains
+// undefined until that tenant is selected for training.
+__attribute__((visibility("default")))
+int32_t qwen36_import_adapter_optimizer_state_host_v1(
+    void* ctx_ptr, int64_t adapter_id,
+    const int64_t* global_layers, const char* const* module_names,
+    void* const* state_ptrs, int64_t slot_count, int64_t optimizer_step
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx && adapter_id > 0 && slot_count >= 0 &&
+                optimizer_step >= 0,
+            "invalid dynamic Adam host import arguments");
+        TORCH_CHECK(!ctx->dynamic_adam_transaction_active,
+            "cannot import dynamic Adam during an active transaction");
+        require_clear_accumulation_for_registry_mutation(ctx);
+        TORCH_CHECK(slot_count == 0 ||
+                (global_layers && module_names && state_ptrs),
+            "dynamic Adam host import is missing slot arrays");
+        size_t adapter_index = 0;
+        TORCH_CHECK(find_canonical_adapter_index(
+                ctx, adapter_id, adapter_index),
+            "dynamic Adam host import adapter not found: ", adapter_id);
+        auto& adapter = ctx->adapters[adapter_index];
+        TORCH_CHECK(adapter.optimizer_step == 0 &&
+                adapter.adam_host_step == -1,
+            "dynamic Adam host import requires a fresh restore adapter: ",
+            adapter_id);
+
+        int64_t expected_slots = 0;
+        std::map<int64_t, std::vector<std::array<at::Tensor, 4>>> snapshot;
+        std::map<int64_t, std::vector<std::array<at::Tensor, 4>>> empty_device;
+        for (const auto& [layer_idx, pairs] : adapter.params) {
+            snapshot.emplace(layer_idx,
+                std::vector<std::array<at::Tensor, 4>>(pairs.size()));
+            empty_device.emplace(layer_idx,
+                std::vector<std::array<at::Tensor, 4>>(pairs.size()));
+            for (const auto& [a, b] : pairs) {
+                TORCH_CHECK(a.requires_grad() == b.requires_grad(),
+                    "dynamic Adam parameter activity mismatch for adapter ",
+                    adapter_id, " layer ", layer_idx);
+                if (a.requires_grad()) ++expected_slots;
+            }
+        }
+        TORCH_CHECK(slot_count == expected_slots,
+            "dynamic Adam host import slot count mismatch for adapter ",
+            adapter_id, ": got ", slot_count,
+            " expected ", expected_slots);
+
+        std::set<std::pair<int64_t, int64_t>> seen;
+        at::NoGradGuard guard;
+        for (int64_t slot = 0; slot < slot_count; ++slot) {
+            TORCH_CHECK(module_names[slot] && module_names[slot][0] != '\0',
+                "dynamic Adam host import has an empty module at slot ", slot);
+            int64_t layer_idx = -1;
+            TORCH_CHECK(global_to_local_layer(
+                    ctx, global_layers[slot], layer_idx),
+                "dynamic Adam host import layer is not owned by this stage: ",
+                global_layers[slot]);
+            const int64_t pair_idx = lora_pair_index(
+                ctx->layer_configs[layer_idx], module_names[slot]);
+            auto params_it = adapter.params.find(layer_idx);
+            TORCH_CHECK(pair_idx >= 0 &&
+                    params_it != adapter.params.end() &&
+                    pair_idx < static_cast<int64_t>(params_it->second.size()),
+                "dynamic Adam host import slot is not in adapter layout: layer=",
+                global_layers[slot], " module=", module_names[slot]);
+            TORCH_CHECK(seen.emplace(layer_idx, pair_idx).second,
+                "duplicate dynamic Adam host import slot: layer=",
+                global_layers[slot], " module=", module_names[slot]);
+            const auto& [a, b] = params_it->second[pair_idx];
+            TORCH_CHECK(a.requires_grad() && b.requires_grad(),
+                "dynamic Adam host import slot is inactive: layer=",
+                global_layers[slot], " module=", module_names[slot]);
+            std::array<at::Tensor, 4> sources;
+            for (int state_index = 0; state_index < 4; ++state_index) {
+                auto* source = reinterpret_cast<at::Tensor*>(
+                    state_ptrs[slot * 4 + state_index]);
+                TORCH_CHECK(source && source->defined() &&
+                        source->device().is_cpu() && source->is_contiguous() &&
+                        source->scalar_type() == at::kFloat,
+                    "dynamic Adam host import requires contiguous CPU FP32 "
+                    "state at slot ", slot, " index ", state_index);
+                sources[state_index] = source->clone();
+            }
+            TORCH_CHECK(sources[0].sizes() == a.sizes() &&
+                    sources[1].sizes() == a.sizes() &&
+                    sources[2].sizes() == b.sizes() &&
+                    sources[3].sizes() == b.sizes(),
+                "dynamic Adam host import shape mismatch at slot ", slot,
+                " layer ", global_layers[slot],
+                " module ", module_names[slot]);
+            snapshot.at(layer_idx)[pair_idx] = std::move(sources);
+        }
+        TORCH_CHECK(static_cast<int64_t>(seen.size()) == expected_slots,
+            "dynamic Adam host import did not cover every active slot");
+        TORCH_CHECK(!ctx->dynamic_adam_transaction_active &&
+                adapter.optimizer_step == 0 &&
+                adapter.adam_host_step == -1,
+            "dynamic Adam changed while preparing host import");
+        adapter.adam_state.swap(empty_device);
+        adapter.adam_host_state.swap(snapshot);
+        adapter.adam_host_step = optimizer_step;
+        adapter.optimizer_step = optimizer_step;
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr,
+            "[q36] import_adapter_optimizer_state_host FAILED: %s\n",
+            e.what());
         return -1;
     }
+}
 
-    // Use PyTorch's compute stream for NCCL — NOT a separate stream.
-    // Separate stream causes "invalid argument" because NCCL communicator
-    // is bound to a CUDA context, and PyTorch's caching allocator stream
-    // may be on a different context. Using the same stream ensures same context.
-    // We store nullptr for stream — moe_forward will use getCurrentCUDAStream(dev).
-    cudaStream_t nccl_stream = nullptr;  // nullptr = use default stream
-
-    // Store as process-level singleton
-    g_nccl_comm = comm;
-    g_nccl_stream = nccl_stream;
-    g_nccl_initialized = true;
-
-    ctx->nccl_comm = comm;
-    ctx->nccl_stream = nccl_stream;
-    ctx->ep_rank = rank;
-    ctx->ep_world_size = world_size;
-
-    // Propagate to layer configs
-    for (auto& lc : ctx->layer_configs) {
-        lc.nccl_comm = (void*)comm;
-        lc.nccl_stream = (void*)nccl_stream;
-    }
-    for (auto& lc : ctx->mtp_layer_configs) {
-        lc.nccl_comm = (void*)comm;
-        lc.nccl_stream = (void*)nccl_stream;
+__attribute__((visibility("default")))
+int64_t qwen36_get_adapter_optimizer_resident_count_v1(
+    void* ctx_ptr, int64_t adapter_id
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx || adapter_id <= 0) return -1;
+    size_t adapter_index = 0;
+    if (!find_canonical_adapter_index(ctx, adapter_id, adapter_index))
+        return -1;
+    int64_t count = 0;
+    for (const auto& [layer_idx, states] :
+         ctx->adapters[adapter_index].adam_state) {
+        (void)layer_idx;
+        for (const auto& state : states) {
+            for (const auto& tensor : state)
+                if (tensor.defined()) ++count;
+        }
     }
-
-    return 0;
+    return count;
 }
 
-// Set NCCL communicator for Expert Parallel all-reduce (legacy, from Rust)
-__attribute__((visibility("default"))) void qwen36_set_nccl_comm(
-    void* ctx_ptr, void* comm_ptr, void* stream_ptr,
-    int32_t ep_rank, int32_t ep_world_size
+__attribute__((visibility("default")))
+int64_t qwen36_get_adapter_optimizer_resident_bytes_v1(
+    void* ctx_ptr, int64_t adapter_id
 ) {
     auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-    ctx->nccl_comm = reinterpret_cast<ncclComm_t>(comm_ptr);
-    ctx->nccl_stream = reinterpret_cast<cudaStream_t>(stream_ptr);
-    ctx->ep_rank = ep_rank;
-    ctx->ep_world_size = ep_world_size;
-    // Propagate NCCL handles to all layer configs so moe_forward can access them
-    for (auto& lc : ctx->layer_configs) {
-        lc.nccl_comm = comm_ptr;
-        lc.nccl_stream = stream_ptr;
-    }
-    for (auto& lc : ctx->mtp_layer_configs) {
-        lc.nccl_comm = comm_ptr;
-        lc.nccl_stream = stream_ptr;
+    if (!ctx || adapter_id <= 0) return -1;
+    size_t adapter_index = 0;
+    if (!find_canonical_adapter_index(ctx, adapter_id, adapter_index))
+        return -1;
+    const uint64_t bytes = dynamic_adam_resident_bytes(
+        ctx->adapters[adapter_index]);
+    return bytes > static_cast<uint64_t>(
+            std::numeric_limits<int64_t>::max())
+        ? std::numeric_limits<int64_t>::max()
+        : static_cast<int64_t>(bytes);
+}
+
+__attribute__((visibility("default")))
+int32_t qwen36_set_adapter_optimizer_tensor(
+    void* ctx_ptr, int64_t adapter_id, int64_t global_layer,
+    const char* module_name, int32_t is_b, int32_t is_v, void* tensor_ptr
+) {
+    try {
+        auto* target = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_optimizer_tensor(
+                ctx_ptr, adapter_id, global_layer, module_name, is_b, is_v));
+        TORCH_CHECK(target && tensor_ptr, "dynamic optimizer tensor not found");
+        auto& source = *reinterpret_cast<at::Tensor*>(tensor_ptr);
+        TORCH_CHECK(source.sizes() == target->sizes(),
+            "dynamic optimizer tensor shape mismatch: expected ", target->sizes(),
+            " got ", source.sizes());
+        at::NoGradGuard guard;
+        target->copy_(source.to(target->device()).to(target->scalar_type()));
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        for (auto& adapter : ctx->adapters) {
+            if (adapter.id != adapter_id) continue;
+            adapter.adam_host_step = -1;
+            break;
+        }
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] set_adapter_optimizer_tensor FAILED: %s\n", e.what());
+        return -1;
     }
 }
 
-// Enable/disable gradient checkpointing
-__attribute__((visibility("default"))) void qwen36_set_checkpoint(void* ctx_ptr, int32_t enable, int64_t group_size) {
+__attribute__((visibility("default")))
+int64_t qwen36_get_adapter_step_count(void* ctx_ptr, int64_t adapter_id) {
     auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-    ctx->use_checkpoint = (enable != 0);
-    ctx->group_size = (group_size > 0) ? group_size : 4;
-    fprintf(stderr, "[q36_ctx] checkpoint: %s, group_size=%ld\n",
-        ctx->use_checkpoint ? "ON" : "OFF", (long)ctx->group_size);
+    if (!ctx || adapter_id <= 0) return -1;
+    for (const auto& adapter : ctx->adapters) {
+        if (adapter.id == adapter_id) return adapter.optimizer_step;
+    }
+    return -1;
 }
 
-// Set attention mask for padding tokens
-__attribute__((visibility("default"), used))
-void qwen36_set_attention_mask(void* ctx_ptr, void* mask_ptr) {
+// Read-only native smoke/debug ABI. Kinds: 0 embedding scatters, 1 column
+// gathers, 2 row reduce-scatters, 3 loss gathers, 4 last local sequence.
+__attribute__((visibility("default")))
+int64_t qwen36_get_sequence_parallel_counter(void* ctx_ptr, int32_t kind) {
     auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-    if (mask_ptr) {
-        ctx->attention_mask = *reinterpret_cast<at::Tensor*>(mask_ptr);
+    if (!ctx || !ctx->sequence_parallel) return -1;
+    switch (kind) {
+        case 0: return ctx->sequence_parallel_embedding_scatter_count;
+        case 1: return ctx->sequence_parallel_all_gather_count;
+        case 2: return ctx->sequence_parallel_reduce_scatter_count;
+        case 3: return ctx->sequence_parallel_loss_gather_count;
+        case 4: return ctx->sequence_parallel_last_local_sequence;
+        default: return -1;
     }
 }
 
-// Utility functions (kept for compatibility)
-__attribute__((visibility("default"))) void* qwen36_gemm(void* a_ptr, void* b_ptr, int transpose_b) {
-    auto& a = *reinterpret_cast<at::Tensor*>(a_ptr);
-    auto& b = *reinterpret_cast<at::Tensor*>(b_ptr);
-    if (transpose_b) return new at::Tensor(at::matmul(a, b.t()));
-    return new at::Tensor(at::matmul(a, b));
+__attribute__((visibility("default")))
+int32_t qwen36_validate_adapter_steps_v1(
+    void* ctx_ptr, const int64_t* adapter_ids,
+    const int64_t* expected_steps, int32_t n_adapters
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx && adapter_ids && expected_steps && n_adapters > 0,
+            "adapter step validation requires a context and non-empty arrays");
+        validate_native_execution_topology_collective(ctx);
+        // This hashes the ordered selection and every tenant's actual clock,
+        // plus the expected clocks, then reduces min/max over TP, EP, and DP
+        // before any rank can return.
+        validate_adapter_collective_registry(
+            ctx, adapter_ids, n_adapters, 0, false, expected_steps);
+        require_canonical_adapter_index(ctx);
+        for (int32_t index = 0; index < n_adapters; ++index) {
+            TORCH_CHECK(adapter_ids[index] > 0 && expected_steps[index] >= 0,
+                "adapter IDs must be positive and expected steps non-negative");
+            size_t adapter_index = 0;
+            TORCH_CHECK(find_canonical_adapter_index(
+                    ctx, adapter_ids[index], adapter_index),
+                "unknown dynamic adapter ID: ", adapter_ids[index]);
+            const auto& adapter = ctx->adapters[adapter_index];
+            TORCH_CHECK(adapter.optimizer_step == expected_steps[index],
+                "adapter ", adapter_ids[index],
+                " optimizer step conflict: expected ", expected_steps[index],
+                ", actual ", adapter.optimizer_step);
+        }
+        return 0;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] validate_adapter_steps FAILED: %s\n", e.what());
+        return -1;
+    }
 }
 
-__attribute__((visibility("default"))) void qwen36_free_tensor(void* tensor_ptr) {
-    if (tensor_ptr) delete reinterpret_cast<at::Tensor*>(tensor_ptr);
+__attribute__((visibility("default")))
+int32_t qwen36_set_adapter_step_count(
+    void* ctx_ptr, int64_t adapter_id, int64_t step_count
+) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    if (!ctx || adapter_id <= 0 || step_count < 0) return -1;
+    for (auto& adapter : ctx->adapters) {
+        if (adapter.id == adapter_id) {
+            if (adapter.optimizer_step != step_count)
+                adapter.adam_host_step = -1;
+            adapter.optimizer_step = step_count;
+            return 0;
+        }
+    }
+    return -1;
 }
 
-// ── Multi-LoRA adapter management ──
+__attribute__((visibility("default")))
+int64_t qwen36_get_lora_count(void* ctx_ptr) {
+    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+    // This legacy accessor is paired with get_lora_a/get_lora_b and therefore
+    // counts only the fixed single-adapter slots. Dynamic adapters have their
+    // own registry and must not make this count exceed those arrays.
+    return (int64_t)ctx->lora_a.size();
+}
 
 __attribute__((visibility("default")))
-int64_t qwen36_add_lora(
-    void* ctx_ptr,
-    int64_t rank, double alpha,
-    const int64_t* target_layers, int64_t num_target_layers,
-    const char* target_modules_str
-) {
+double qwen36_eval_step(void* ctx_ptr, void* input_ids_ptr, void* target_mask_ptr, void* attention_mask_ptr) {
     try {
         auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-        TrainingContext::LoRAAdapter adapter;
-        adapter.id = ++ctx->next_adapter_id;
-        adapter.rank = rank;
-        adapter.alpha = alpha;
-        if (target_layers && num_target_layers > 0) {
-            for (int64_t i = 0; i < num_target_layers; i++)
-                adapter.target_layers.insert(target_layers[i]);
+        struct EvalCacheGuard {
+            TrainingContext* ctx;
+            ~EvalCacheGuard() {
+                // Evaluation builds projection caches under no-grad. They
+                // must never be reused by the following training step.
+                if (!ctx) return;
+                ctx->lora_cache_valid = false;
+                ctx->lora_batch_valid = false;
+            }
+        } cache_guard{ctx};
+        validate_native_execution_topology_collective(ctx);
+        auto* input_ids_tensor = reinterpret_cast<at::Tensor*>(input_ids_ptr);
+        auto* target_mask_tensor = reinterpret_cast<at::Tensor*>(target_mask_ptr);
+        auto* attention_mask_tensor = attention_mask_ptr
+            ? reinterpret_cast<at::Tensor*>(attention_mask_ptr)
+            : nullptr;
+        bool local_preflight_valid = input_ids_tensor && target_mask_tensor &&
+            input_ids_tensor->is_cuda() && target_mask_tensor->is_cuda() &&
+            input_ids_tensor->device() == target_mask_tensor->device() &&
+            input_ids_tensor->dim() == 2 && target_mask_tensor->dim() == 2 &&
+            input_ids_tensor->size(1) > 1 &&
+            input_ids_tensor->sizes() == target_mask_tensor->sizes() &&
+            input_ids_tensor->scalar_type() == at::kLong &&
+            supported_mask_dtype(*target_mask_tensor) &&
+            ctx->adapters.empty();
+        if (local_preflight_valid && attention_mask_tensor) {
+            local_preflight_valid = attention_mask_tensor->is_cuda() &&
+                attention_mask_tensor->device() == input_ids_tensor->device() &&
+                attention_mask_tensor->dim() == 2 &&
+                attention_mask_tensor->sizes() == input_ids_tensor->sizes() &&
+                supported_mask_dtype(*attention_mask_tensor);
         }
-        if (target_modules_str) {
-            std::string s(target_modules_str);
-            std::stringstream ss(s);
-            std::string item;
-            while (std::getline(ss, item, ','))
-                adapter.target_modules.insert(item);
+        if (local_preflight_valid &&
+                (ctx->nccl_comm || ctx->dp_comm || ctx->tp_comm ||
+                    ctx->cp_comm)) {
+            local_preflight_valid =
+                input_ids_tensor->device().index() == ctx->cuda_device;
         }
-        for (int64_t i = 0; i < ctx->num_layers; i++) {
-            if (!adapter.target_layers.empty() && adapter.target_layers.find(i) == adapter.target_layers.end())
-                continue;
-            int64_t w_offset = 0;
-            for (int64_t j = 0; j < i; j++)
-                w_offset += weight_count_for_layer(ctx->layer_configs[j]);
-            int64_t num_pairs;
-            const int64_t* proj_indices;
-            if (ctx->layer_configs[i].layer_type == 0) {
-                static const int64_t full_indices[] = {2, 4, 6, 7};
-                proj_indices = full_indices;
-                num_pairs = 4;
-            } else {
-                static const int64_t linear_indices[] = {2, 3, 10};
-                proj_indices = linear_indices;
-                num_pairs = 3;
-            }
-            std::vector<std::pair<at::Tensor, at::Tensor>> pairs;
-            std::vector<std::array<at::Tensor, 4>> adam_states;
-            for (int k = 0; k < num_pairs; k++) {
-                auto* base = ctx->weight_ptrs[w_offset + proj_indices[k]];
-                int64_t out_f = base->size(0), in_f = base->size(1);
-                auto a = at::randn({rank, in_f}, at::TensorOptions().dtype(ctx->compute_type).device(base->device())) * 0.01;
-                auto b = at::zeros({out_f, rank}, at::TensorOptions().dtype(ctx->compute_type).device(base->device()));
-                a.set_requires_grad(true);
-                b.set_requires_grad(true);
-                // Adam state: FP32 for numerical stability
-                auto opts_f32 = at::TensorOptions().dtype(at::kFloat).device(base->device());
-                adam_states.push_back({
-                    at::zeros(a.sizes(), opts_f32), at::zeros(a.sizes(), opts_f32),
-                    at::zeros(b.sizes(), opts_f32), at::zeros(b.sizes(), opts_f32)
-                });
-                pairs.emplace_back(std::move(a), std::move(b));
-            }
-            adapter.params[i] = std::move(pairs);
-            adapter.adam_state[i] = std::move(adam_states);
+        TORCH_CHECK(adapter_collective_all_succeeded(
+                ctx, local_preflight_valid) && local_preflight_valid,
+            "native Qwen fixed-LoRA eval requires matching CUDA input and "
+            "mask tensors on every distributed rank and no dynamic adapters");
+        auto& input_ids = *input_ids_tensor;
+        auto& target_mask = *target_mask_tensor;
+        if (!ctx->nccl_comm && !ctx->dp_comm && !ctx->tp_comm &&
+                !ctx->cp_comm)
+            ctx->cuda_device = input_ids.device().index();
+        c10::cuda::set_device(ctx->cuda_device);
+        cudaSetDevice(ctx->cuda_device);
+        validate_fixed_collective_registry(ctx);
+        const bool replica_inputs_match = replica_input_signatures_match(
+            ctx, input_ids, target_mask, attention_mask_tensor);
+        TORCH_CHECK(adapter_collective_all_succeeded(
+                ctx, replica_inputs_match) && replica_inputs_match,
+            "native Qwen fixed-LoRA eval input shape, dtype, or value differs "
+            "across TP, CP, or replicated EP ranks");
+        if (attention_mask_tensor) {
+            auto& attention_mask = *attention_mask_tensor;
+            validate_linear_attention_mask(ctx, attention_mask);
+            ctx->attention_mask = attention_mask;
+        } else {
+            ctx->attention_mask = at::Tensor();
+            ctx->attention_lengths = at::Tensor();
         }
-        int64_t id = adapter.id;
-        ctx->adapters.push_back(std::move(adapter));
-        ctx->lora_cache_valid = false;
-        ctx->lora_batch_valid = false;
-        fprintf(stderr, "[q36_lora] added adapter %ld: rank=%ld alpha=%.1f\n", (long)id, (long)rank, alpha);
-        return id;
+        elide_trivial_attention_mask(ctx);
+        at::AutoGradMode no_grad(false);
+        auto hidden = ctx->use_checkpoint ? forward_full_checkpoint(ctx, input_ids) : forward_full(ctx, input_ids);
+        auto loss_hidden = sequence_parallel_loss_gather(ctx, hidden);
+        auto loss = compute_loss_fused(
+            ctx, loss_hidden, input_ids, target_mask, ctx->vocab_size);
+        return loss.item<double>();
     } catch (const std::exception& e) {
-        fprintf(stderr, "[q36] add_lora FAILED: %s\n", e.what());
-        return -1;
+        fprintf(stderr, "[q36] eval_step FAILED: %s\n", e.what());
+        return -1.0;
     }
 }
 
-__attribute__((visibility("default")))
-int32_t qwen36_remove_lora(void* ctx_ptr, int64_t adapter_id) {
+// Evaluate selected dynamic tenants independently. The registry is scoped to
+// one tenant at a time so heterogeneous ranks/modules cannot cross-contaminate
+// routing or activation-level LoRA batches. Loss numerators/counts are reduced
+// over source-sharded EP and expert-DP just like selected training reports.
+__attribute__((visibility("default"))) int32_t qwen36_eval_multi_lora_selected_v1(
+    void* ctx_ptr,
+    void* input_ids_ptr,
+    void* target_mask_ptr,
+    void* attention_mask_ptr,
+    const int64_t* adapter_ids,
+    int32_t n_adapters,
+    double* adapter_losses_out,
+    int32_t adapter_loss_capacity
+) {
     auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-    for (auto it = ctx->adapters.begin(); it != ctx->adapters.end(); ++it) {
-        if (it->id == adapter_id) {
-            ctx->adapters.erase(it);
+    std::vector<TrainingContext::LoRAAdapter> original;
+    std::vector<TrainingContext::LoRAAdapter> selected;
+    std::vector<size_t> selected_indexes;
+    bool registry_detached = false;
+    bool adapter_id_index_was_valid = false;
+    auto restore_registry = [&]() {
+        if (!ctx || !registry_detached) return;
+        for (size_t i = 0; i < selected.size(); ++i) {
+            original[selected_indexes[i]] = std::move(selected[i]);
+        }
+        ctx->adapters.clear();
+        ctx->adapters.swap(original);
+        ctx->adapter_id_index_valid = adapter_id_index_was_valid;
+        registry_detached = false;
+    };
+    try {
+        TORCH_CHECK(ctx && adapter_ids && n_adapters > 0 &&
+                adapter_losses_out && adapter_loss_capacity >= n_adapters,
+            "selected multi-LoRA eval requires valid IDs and output storage");
+        validate_dynamic_context_health(ctx);
+        validate_native_execution_topology_collective(ctx);
+        TORCH_CHECK(ctx->pp_world_size == 1 && ctx->cp_world_size == 1,
+            "selected multi-LoRA eval requires PP=1 and CP=1");
+        auto* input = reinterpret_cast<at::Tensor*>(input_ids_ptr);
+        auto* target = reinterpret_cast<at::Tensor*>(target_mask_ptr);
+        auto* attention = attention_mask_ptr
+            ? reinterpret_cast<at::Tensor*>(attention_mask_ptr) : nullptr;
+        TORCH_CHECK(input && target && input->is_cuda() && target->is_cuda() &&
+                input->dim() == 2 && target->sizes() == input->sizes() &&
+                (input->size(0) == 1 || input->size(0) == n_adapters),
+            "selected multi-LoRA eval input batch must be [1, seq] or [n, seq]");
+        TORCH_CHECK(input->scalar_type() == at::kLong &&
+                supported_mask_dtype(*target),
+            "selected multi-LoRA eval input dtypes are invalid");
+        if (attention) {
+            TORCH_CHECK(attention->is_cuda() && attention->sizes() == input->sizes() &&
+                    supported_mask_dtype(*attention),
+                "selected multi-LoRA eval attention mask is invalid");
+        }
+        validate_adapter_collective_registry(
+            ctx, adapter_ids, n_adapters, 0, false);
+
+        require_canonical_adapter_index(ctx);
+        selected.reserve(n_adapters);
+        selected_indexes.reserve(n_adapters);
+        std::vector<uint8_t> requested(ctx->adapters.size(), 0);
+        for (int32_t i = 0; i < n_adapters; ++i) {
+            TORCH_CHECK(adapter_ids[i] > 0,
+                "selected eval adapter IDs must be positive");
+            size_t index = 0;
+            TORCH_CHECK(find_canonical_adapter_index(
+                    ctx, adapter_ids[i], index),
+                "unknown selected eval adapter ID: ", adapter_ids[i]);
+            TORCH_CHECK(!requested[index],
+                "duplicate selected eval adapter ID: ", adapter_ids[i]);
+            requested[index] = 1;
+            selected_indexes.push_back(index);
+        }
+        adapter_id_index_was_valid = ctx->adapter_id_index_valid;
+        ctx->adapter_id_index_valid = false;
+        original.swap(ctx->adapters);
+        registry_detached = true;
+        for (int32_t i = 0; i < n_adapters; ++i) {
+            selected.push_back(std::move(
+                original[selected_indexes[static_cast<size_t>(i)]]));
+        }
+
+        const at::Tensor saved_attention_mask = ctx->attention_mask;
+        const at::Tensor saved_attention_lengths = ctx->attention_lengths;
+        struct EvalMaskGuard {
+            TrainingContext* ctx;
+            at::Tensor mask;
+            at::Tensor lengths;
+            ~EvalMaskGuard() {
+                ctx->attention_mask = mask;
+                ctx->attention_lengths = lengths;
+            }
+        } mask_guard{ctx, saved_attention_mask, saved_attention_lengths};
+        for (int32_t i = 0; i < n_adapters; ++i) {
+            ctx->adapters.clear();
+            ctx->adapters.push_back(selected[i]);
             ctx->lora_cache_valid = false;
             ctx->lora_batch_valid = false;
-            return 1;
+            const bool shared_row = input->size(0) == 1;
+            auto ids = shared_row ? *input : input->narrow(0, i, 1).contiguous();
+            auto targets = shared_row ? *target : target->narrow(0, i, 1).contiguous();
+            if (attention) {
+                ctx->attention_mask = shared_row
+                    ? *attention : attention->narrow(0, i, 1).contiguous();
+                elide_trivial_attention_mask(ctx);
+            } else {
+                ctx->attention_mask = at::Tensor();
+                ctx->attention_lengths = at::Tensor();
+            }
+            at::AutoGradMode no_grad(false);
+            auto hidden = forward_full(ctx, ids);
+            auto loss = compute_loss_fused(
+                ctx, hidden, ids, targets, ctx->vocab_size);
+            auto counts = targets.narrow(1, 1, targets.size(1) - 1)
+                .to(at::kFloat).sum().reshape({1});
+            auto numerators = (loss * counts).reshape({1});
+            auto global_loss = global_dynamic_adapter_losses(
+                ctx, numerators, counts);
+            adapter_losses_out[i] = global_loss.item<double>();
         }
+        restore_registry();
+        return 0;
+    } catch (const std::exception& error) {
+        try { restore_registry(); } catch (...) {}
+        fprintf(stderr, "[q36] selected multi-LoRA eval FAILED: %s\n",
+            error.what());
+        return -1;
+    } catch (...) {
+        try { restore_registry(); } catch (...) {}
+        fprintf(stderr, "[q36] selected multi-LoRA eval FAILED: unknown exception\n");
+        return -1;
     }
-    return 0;
+}
+
+static at::Tensor qwen36_copy_host_i64_batch(
+    TrainingContext* ctx,
+    const int64_t* data,
+    int64_t batch_size,
+    int64_t seq_len,
+    const char* name,
+    at::Tensor* pinned_staging
+) {
+    TORCH_CHECK(ctx, name, " requires a valid training context");
+    TORCH_CHECK(data, name, " requires a non-null host pointer");
+    TORCH_CHECK(batch_size > 0 && seq_len > 0,
+        name, " requires positive batch_size and seq_len");
+    TORCH_CHECK(batch_size <= std::numeric_limits<int64_t>::max() / seq_len,
+        name, " batch shape overflows int64");
+    c10::cuda::set_device(ctx->cuda_device);
+    cudaSetDevice(ctx->cuda_device);
+    auto host = at::from_blob(
+        const_cast<int64_t*>(data), {batch_size, seq_len},
+        at::TensorOptions().device(at::kCPU).dtype(at::kLong));
+    if (env_enabled("QWEN36_HOST_PINNED_STAGING")) {
+        TORCH_CHECK(pinned_staging,
+            name, " pinned staging destination is unavailable");
+        if (!pinned_staging->defined() || pinned_staging->dim() != 2 ||
+                pinned_staging->size(0) != batch_size ||
+                pinned_staging->size(1) != seq_len) {
+            *pinned_staging = at::empty(
+                {batch_size, seq_len},
+                at::TensorOptions().device(at::kCPU).dtype(at::kLong)
+                    .pinned_memory(true));
+        }
+        // This CPU copy is intentionally synchronous: the source is a
+        // pageable IPC slab. Only the subsequent pinned H2D transfer is
+        // asynchronous; the host-i64 ABI still returns after its scalar
+        // result is materialized, so this does not imply request overlap.
+        pinned_staging->copy_(host);
+        return pinned_staging->to(
+            at::Device(at::kCUDA, ctx->cuda_device), at::kLong,
+            /*non_blocking=*/true, /*copy=*/true);
+    }
+    // Shared-memory storage remains owned by the IPC channel. A blocking H2D
+    // copy makes the returned CUDA tensor independent before the worker posts
+    // completion and permits the coordinator to reuse the slab.
+    return host.to(at::Device(at::kCUDA, ctx->cuda_device), at::kLong,
+        /*non_blocking=*/false, /*copy=*/true);
+}
+
+struct Qwen36HostI64Batch {
+    at::Tensor input;
+    at::Tensor targets;
+    at::Tensor attention;
+};
+
+static Qwen36HostI64Batch qwen36_copy_host_i64_batch_collective(
+    TrainingContext* ctx,
+    const int64_t* input_ids,
+    const int64_t* target_mask,
+    const int64_t* attention_mask,
+    int64_t batch_size,
+    int64_t seq_len,
+    bool local_request_valid,
+    const char* name
+) {
+    const bool shape_valid = batch_size > 0 && seq_len > 0 &&
+        batch_size <= std::numeric_limits<int64_t>::max() / seq_len;
+    const bool local_preflight_valid = ctx && input_ids && target_mask &&
+        attention_mask && shape_valid && local_request_valid;
+    const bool globally_valid = adapter_collective_all_succeeded(
+        ctx, local_preflight_valid);
+    TORCH_CHECK(local_preflight_valid && globally_valid,
+        name, " request pointers, dimensions, or distributed metadata are "
+        "invalid or differ across ranks");
+
+    Qwen36HostI64Batch batch;
+    std::string local_copy_error;
+    try {
+        batch.input = qwen36_copy_host_i64_batch(
+            ctx, input_ids, batch_size, seq_len, name,
+            &ctx->host_i64_input_staging);
+        batch.targets = qwen36_copy_host_i64_batch(
+            ctx, target_mask, batch_size, seq_len, name,
+            &ctx->host_i64_target_staging);
+        batch.attention = qwen36_copy_host_i64_batch(
+            ctx, attention_mask, batch_size, seq_len, name,
+            &ctx->host_i64_attention_staging);
+    } catch (const std::exception& error) {
+        local_copy_error = error.what();
+    }
+    const bool local_copy_succeeded = local_copy_error.empty();
+    const bool globally_copied = adapter_collective_all_succeeded(
+        ctx, local_copy_succeeded);
+    TORCH_CHECK(local_copy_succeeded && globally_copied,
+        name, " H2D copy failed on at least one distributed rank",
+        local_copy_error.empty() ? "" : ": ", local_copy_error);
+    return batch;
 }
 
 __attribute__((visibility("default")))
-int64_t qwen36_list_lora(void* ctx_ptr, int64_t* out_ids, int64_t max_count) {
-    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-    int64_t count = (int64_t)ctx->adapters.size();
-    if (count > max_count) count = max_count;
-    for (int64_t i = 0; i < count; i++)
-        out_ids[i] = ctx->adapters[i].id;
-    return count;
+double qwen36_train_step_host_i64(
+    void* ctx_ptr,
+    const int64_t* input_ids,
+    const int64_t* target_mask,
+    const int64_t* attention_mask,
+    int64_t batch_size,
+    int64_t seq_len
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        validate_native_execution_topology_collective(ctx);
+        auto batch = qwen36_copy_host_i64_batch_collective(
+            ctx, input_ids, target_mask, attention_mask,
+            batch_size, seq_len, true, "host fixed-LoRA train");
+        return qwen36_train_step(
+            ctx, &batch.input, &batch.targets, &batch.attention);
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] train_step_host_i64 FAILED: %s\n", e.what());
+        return -1.0;
+    }
 }
 
 __attribute__((visibility("default")))
-int64_t qwen36_get_lora_count(void* ctx_ptr) {
-    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-    int64_t total = (int64_t)ctx->lora_a.size();
-    for (auto& adapter : ctx->adapters)
-        for (auto& [layer_idx, pairs] : adapter.params)
-            total += (int64_t)pairs.size() * 2;
-    return total;
+double qwen36_train_multi_lora_host_i64(
+    void* ctx_ptr,
+    const int64_t* input_ids,
+    const int64_t* target_mask,
+    const int64_t* attention_mask,
+    int64_t batch_size,
+    int64_t seq_len,
+    int32_t n_total,
+    int32_t lora_rank,
+    const int64_t* adapter_ids,
+    int32_t n_adapter_ids
+) {
+    try {
+        (void)lora_rank;  // Retained for the ABI22 host wire contract.
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        validate_native_execution_topology_collective(ctx);
+        const bool local_request_valid = n_total > 0 &&
+            n_adapter_ids >= 0 &&
+            (n_adapter_ids == 0 || n_adapter_ids == n_total) &&
+            (n_adapter_ids == 0 || adapter_ids) &&
+            (n_adapter_ids != 0 ||
+                static_cast<int32_t>(ctx->adapters.size()) == n_total);
+        auto batch = qwen36_copy_host_i64_batch_collective(
+            ctx, input_ids, target_mask, attention_mask,
+            batch_size, seq_len, local_request_valid, "host multi-LoRA train");
+        if (n_adapter_ids == 0) {
+            std::vector<int64_t> all_adapter_ids;
+            all_adapter_ids.reserve(ctx->adapters.size());
+            for (const auto& adapter : ctx->adapters)
+                all_adapter_ids.push_back(adapter.id);
+            return qwen36_train_multi_lora_selected_v2(
+                ctx, &batch.input, &batch.targets, &batch.attention,
+                all_adapter_ids.data(), n_total);
+        }
+        return qwen36_train_multi_lora_selected_v2(
+            ctx, &batch.input, &batch.targets, &batch.attention,
+            adapter_ids, n_adapter_ids);
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] train_multi_lora_host_i64 FAILED: %s\n", e.what());
+        return -1.0;
+    }
 }
 
 __attribute__((visibility("default")))
-double qwen36_eval_step(void* ctx_ptr, void* input_ids_ptr, void* target_mask_ptr, void* attention_mask_ptr) {
+int32_t qwen36_train_multi_lora_host_i64_v2(
+    void* ctx_ptr,
+    const int64_t* input_ids,
+    const int64_t* target_mask,
+    const int64_t* attention_mask,
+    int64_t batch_size,
+    int64_t seq_len,
+    int32_t n_total,
+    int32_t lora_rank,
+    const int64_t* adapter_ids,
+    int32_t n_adapter_ids,
+    double* aggregate_loss_out,
+    double* adapter_losses_out,
+    int32_t adapter_loss_capacity
+) {
     try {
+        (void)lora_rank;
         auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-        auto& input_ids = *reinterpret_cast<at::Tensor*>(input_ids_ptr);
-        auto& target_mask = *reinterpret_cast<at::Tensor*>(target_mask_ptr);
-        if (attention_mask_ptr)
-            ctx->attention_mask = *reinterpret_cast<at::Tensor*>(attention_mask_ptr);
-        at::AutoGradMode no_grad(false);
-        auto hidden = ctx->use_checkpoint ? forward_full_checkpoint(ctx, input_ids) : forward_full(ctx, input_ids);
-        auto loss = compute_loss(ctx, hidden, input_ids, target_mask, ctx->vocab_size);
-        return loss.item<double>();
+        validate_native_execution_topology_collective(ctx);
+        const bool report_requested = aggregate_loss_out || adapter_losses_out;
+        const bool report_valid =
+            (!report_requested && !aggregate_loss_out && !adapter_losses_out) ||
+            (aggregate_loss_out && adapter_losses_out &&
+                adapter_loss_capacity >= n_adapter_ids);
+        const bool local_request_valid = n_total > 0 && adapter_ids &&
+            n_adapter_ids > 0 && n_adapter_ids == n_total &&
+            report_valid;
+        auto batch = qwen36_copy_host_i64_batch_collective(
+            ctx, input_ids, target_mask, attention_mask,
+            batch_size, seq_len, local_request_valid,
+            "host multi-LoRA report train");
+        return qwen36_train_multi_lora_selected_v3(
+            ctx, &batch.input, &batch.targets, &batch.attention,
+            adapter_ids, n_adapter_ids,
+            aggregate_loss_out, adapter_losses_out, adapter_loss_capacity);
     } catch (const std::exception& e) {
-        fprintf(stderr, "[q36] eval_step FAILED: %s\n", e.what());
+        fprintf(stderr,
+            "[q36] train_multi_lora_host_i64_v2 FAILED: %s\n", e.what());
+        return -1;
+    }
+}
+
+__attribute__((visibility("default")))
+double qwen36_eval_step_host_i64(
+    void* ctx_ptr,
+    const int64_t* input_ids,
+    const int64_t* target_mask,
+    const int64_t* attention_mask,
+    int64_t batch_size,
+    int64_t seq_len
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        validate_native_execution_topology_collective(ctx);
+        auto batch = qwen36_copy_host_i64_batch_collective(
+            ctx, input_ids, target_mask, attention_mask,
+            batch_size, seq_len, true, "host fixed-LoRA eval");
+        return qwen36_eval_step(
+            ctx, &batch.input, &batch.targets, &batch.attention);
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] eval_step_host_i64 FAILED: %s\n", e.what());
         return -1.0;
     }
 }
 
+__attribute__((visibility("default"))) int32_t qwen36_eval_multi_lora_host_i64_v1(
+    void* ctx_ptr,
+    const int64_t* input_ids,
+    const int64_t* target_mask,
+    const int64_t* attention_mask,
+    int64_t batch_size,
+    int64_t seq_len,
+    const int64_t* adapter_ids,
+    int32_t n_adapters,
+    double* adapter_losses_out,
+    int32_t adapter_loss_capacity
+) {
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        validate_native_execution_topology_collective(ctx);
+        const bool local_request_valid = adapter_ids && n_adapters > 0 &&
+            adapter_losses_out && adapter_loss_capacity >= n_adapters;
+        auto batch = qwen36_copy_host_i64_batch_collective(
+            ctx, input_ids, target_mask, attention_mask,
+            batch_size, seq_len, local_request_valid,
+            "host selected multi-LoRA eval");
+        return qwen36_eval_multi_lora_selected_v1(
+            ctx, &batch.input, &batch.targets, &batch.attention,
+            adapter_ids, n_adapters,
+            adapter_losses_out, adapter_loss_capacity);
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] selected eval host_i64 FAILED: %s\n", e.what());
+        return -1;
+    }
+}
+
 __attribute__((visibility("default")))
 int64_t qwen36_get_step_count(void* ctx_ptr) {
-    return (int64_t)reinterpret_cast<TrainingContext*>(ctx_ptr)->step_count;
+    return (int64_t)reinterpret_cast<TrainingContext*>(ctx_ptr)->fixed_optimizer_step;
+}
+
+__attribute__((visibility("default")))
+int64_t qwen36_get_dynamic_finalizer_count(void* ctx_ptr) {
+    if (!ctx_ptr) return -1;
+    return reinterpret_cast<TrainingContext*>(ctx_ptr)->dynamic_finalizer_count;
+}
+
+__attribute__((visibility("default")))
+int64_t qwen36_get_dynamic_adam_launch_count(void* ctx_ptr) {
+    if (!ctx_ptr) return -1;
+    return reinterpret_cast<TrainingContext*>(ctx_ptr)->dynamic_adam_launch_count;
+}
+
+__attribute__((visibility("default")))
+int64_t qwen36_get_dynamic_train_batch_count(void* ctx_ptr) {
+    if (!ctx_ptr) return -1;
+    return reinterpret_cast<TrainingContext*>(
+        ctx_ptr)->multi_lora_invocation;
+}
+
+__attribute__((visibility("default")))
+int64_t qwen36_get_lora_batch_projection_build_count(void* ctx_ptr) {
+    if (!ctx_ptr) return -1;
+    return reinterpret_cast<TrainingContext*>(
+        ctx_ptr)->lora_batch_projection_build_count;
+}
+
+__attribute__((visibility("default")))
+int64_t qwen36_get_lora_batch_scaling_upload_count(void* ctx_ptr) {
+    if (!ctx_ptr) return -1;
+    return reinterpret_cast<TrainingContext*>(
+        ctx_ptr)->lora_batch_scaling_upload_count;
+}
+
+__attribute__((visibility("default")))
+int32_t qwen36_get_accumulation_active(void* ctx_ptr) {
+    if (!ctx_ptr) return -1;
+    return reinterpret_cast<TrainingContext*>(ctx_ptr)->accumulation_active
+        ? 1 : 0;
+}
+
+// 0 means usable, -1 means null or quarantined. The worker uses this after a
+// failed dynamic-LoRA transaction to decide whether the session can continue.
+__attribute__((visibility("default")))
+int32_t qwen36_get_context_health(void* ctx_ptr) {
+    if (!ctx_ptr) return -1;
+    return reinterpret_cast<TrainingContext*>(ctx_ptr)->poisoned ? -1 : 0;
+}
+
+__attribute__((visibility("default")))
+double qwen36_get_accumulated_token_weight(void* ctx_ptr) {
+    if (!ctx_ptr) return -1.0;
+    return reinterpret_cast<TrainingContext*>(
+        ctx_ptr)->accumulated_token_weight;
+}
+
+// Restore the Adam bias-correction clock independently from tensor state.
+// Checkpoint loading imports m/v through a separate ABI, so omitting this
+// value would resume the next update as step 1 even for a mature optimizer.
+__attribute__((visibility("default")))
+int32_t qwen36_set_step_count(void* ctx_ptr, int64_t step_count) {
+    if (!ctx_ptr || step_count < 0) return -1;
+    reinterpret_cast<TrainingContext*>(ctx_ptr)->fixed_optimizer_step = step_count;
+    return 0;
 }
 
 __attribute__((visibility("default")))
@@ -3670,19 +17674,37 @@ int64_t qwen36_export_optimizer_state(void* ctx_ptr, void** m_ptrs, void** v_ptr
 
 __attribute__((visibility("default")))
 int64_t qwen36_import_optimizer_state(void* ctx_ptr, void** m_ptrs, void** v_ptrs, int64_t count) {
-    auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
-    int64_t imported = 0;
-    for (int64_t i = 0; i < count && i < (int64_t)ctx->adam_m.size(); i++) {
-        auto* src_m = reinterpret_cast<at::Tensor*>(m_ptrs[i]);
-        auto* src_v = reinterpret_cast<at::Tensor*>(v_ptrs[i]);
-        if (src_m && src_v) {
-            ctx->adam_m[i] = src_m->clone();
-            ctx->adam_v[i] = src_v->clone();
-            imported++;
+    try {
+        auto* ctx = reinterpret_cast<TrainingContext*>(ctx_ptr);
+        TORCH_CHECK(ctx, "null training context");
+        TORCH_CHECK(count >= 0, "negative optimizer state count");
+        TORCH_CHECK(count <= (int64_t)ctx->adam_m.size() &&
+                count <= (int64_t)ctx->adam_v.size(),
+            "optimizer state count exceeds native state: count=", count,
+            " m=", ctx->adam_m.size(), " v=", ctx->adam_v.size());
+        TORCH_CHECK(count == 0 || (m_ptrs && v_ptrs),
+            "null optimizer state pointer array");
+        at::NoGradGuard guard;
+        for (int64_t i = 0; i < count; i++) {
+            auto* src_m = reinterpret_cast<at::Tensor*>(m_ptrs[i]);
+            auto* src_v = reinterpret_cast<at::Tensor*>(v_ptrs[i]);
+            TORCH_CHECK(src_m && src_v, "null optimizer tensor at index ", i);
+            auto& target_m = ctx->adam_m[i];
+            auto& target_v = ctx->adam_v[i];
+            TORCH_CHECK(src_m->sizes() == target_m.sizes(),
+                "Adam m shape mismatch at index ", i,
+                ": expected ", target_m.sizes(), " got ", src_m->sizes());
+            TORCH_CHECK(src_v->sizes() == target_v.sizes(),
+                "Adam v shape mismatch at index ", i,
+                ": expected ", target_v.sizes(), " got ", src_v->sizes());
+            target_m.copy_(src_m->to(target_m.device()).to(target_m.scalar_type()));
+            target_v.copy_(src_v->to(target_v.device()).to(target_v.scalar_type()));
         }
+        return count;
+    } catch (const std::exception& e) {
+        fprintf(stderr, "[q36] import_optimizer_state FAILED: %s\n", e.what());
+        return -1;
     }
-    return imported;
 }
 
-
 }  // extern "C"
diff --git a/crates/rustrain-qwen3-6/src/checkpoint.rs b/crates/rustrain-qwen3-6/src/checkpoint.rs
new file mode 100644
index 00000000..78d90acc
--- /dev/null
+++ b/crates/rustrain-qwen3-6/src/checkpoint.rs
@@ -0,0 +1,5191 @@
+//! Qwen checkpoint save/load: LoRA parameters, Adam state, and parallel metadata.
+
+use anyhow::{bail, Context, Result};
+use rustrain_parallel::topology::{ParallelAxis, ParallelTopology, RankCoordinates};
+use serde::{Deserialize, Serialize};
+use std::collections::BTreeMap;
+use std::env;
+use std::io::Read;
+use std::path::{Path, PathBuf};
+use tch::Tensor;
+
+use crate::config::Qwen36RuntimeConfig;
+use crate::lora::{
+    native_lora_slots, Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule,
+};
+
+const TP_CHECKPOINT_FORMAT: &str = "rustrain-checkpoint-v5-parallel";
+const STAGE_UNION_CHECKPOINT_FORMAT: &str = "rustrain-checkpoint-v6-stage-union";
+const PROJECTION_AWARE_TP_CHECKPOINT_FORMAT: &str = "rustrain-checkpoint-v4-tp";
+const LEGACY_TP_CHECKPOINT_FORMAT: &str = "rustrain-checkpoint-v3-tp";
+const RANK_RECEIPT_FORMAT: &str = "rustrain-checkpoint-rank-receipt-v1";
+const RANK_RECEIPT_FILE: &str = "rank-receipt.json";
+const RANK_RECEIPT_VERSION: u32 = 1;
+
+pub fn is_legacy_tensor_parallel_checkpoint(manifest: &CheckpointManifest) -> bool {
+    manifest.format == LEGACY_TP_CHECKPOINT_FORMAT
+}
+
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum LoraTpShardLayout {
+    #[default]
+    LatentRank,
+    ColumnParallel,
+    /// Column-parallel projection whose global rows use flat
+    /// `[Q_all | K_all | V_all]` storage while each rank stores packed
+    /// `[Q_local | K_local | V_local]` rows.
+    FlatQkvColumnParallel {
+        q_rows: i64,
+        k_rows: i64,
+        v_rows: i64,
+    },
+    RowParallel,
+    /// EP shards experts on axis 0 while TP stores packed local
+    /// `[gate_local | up_local]` rows on the projection axis.
+    RoutedExpertFusedGateUp,
+    /// EP shards experts on axis 0 while TP shards the input projection axis.
+    RoutedExpertDown,
+}
+
+pub fn lora_tp_shard_layout(
+    module: Qwen36LoraTargetModule,
+    config: &Qwen36RuntimeConfig,
+) -> LoraTpShardLayout {
+    match module {
+        Qwen36LoraTargetModule::QProj
+        | Qwen36LoraTargetModule::KProj
+        | Qwen36LoraTargetModule::VProj
+        | Qwen36LoraTargetModule::InProjZ
+        | Qwen36LoraTargetModule::InProjA
+        | Qwen36LoraTargetModule::InProjB
+        | Qwen36LoraTargetModule::GateProj
+        | Qwen36LoraTargetModule::UpProj
+        | Qwen36LoraTargetModule::SharedGateProj
+        | Qwen36LoraTargetModule::SharedUpProj => LoraTpShardLayout::ColumnParallel,
+        Qwen36LoraTargetModule::InProjQkv => LoraTpShardLayout::FlatQkvColumnParallel {
+            q_rows: config.linear_num_key_heads * config.linear_key_head_dim,
+            k_rows: config.linear_num_key_heads * config.linear_key_head_dim,
+            v_rows: config.linear_num_value_heads * config.linear_value_head_dim,
+        },
+        Qwen36LoraTargetModule::OProj
+        | Qwen36LoraTargetModule::OutProj
+        | Qwen36LoraTargetModule::DownProj
+        | Qwen36LoraTargetModule::SharedDownProj => LoraTpShardLayout::RowParallel,
+        Qwen36LoraTargetModule::ExpertsGateUpProj => LoraTpShardLayout::RoutedExpertFusedGateUp,
+        Qwen36LoraTargetModule::ExpertsDownProj => LoraTpShardLayout::RoutedExpertDown,
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct LoraSlotIdentity {
+    pub index: usize,
+    pub layer: usize,
+    pub module: String,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct PipelineStageCheckpointManifest {
+    pub pipeline_rank: usize,
+    pub pipeline_size: usize,
+    pub global_num_layers: usize,
+    pub layer_start: usize,
+    pub layer_end: usize,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct StageUnionCheckpointMetadata {
+    pub pipeline_stage: PipelineStageCheckpointManifest,
+    pub fixed_target_layers: Vec<usize>,
+    pub fixed_target_modules: Vec<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ParallelCheckpointManifest {
+    pub world_size: usize,
+    pub tensor_model_parallel_size: usize,
+    pub pipeline_model_parallel_size: usize,
+    pub data_parallel_size: usize,
+    pub expert_model_parallel_size: usize,
+    pub context_parallel_size: usize,
+    pub global_rank: usize,
+    pub tensor_model_parallel_rank: usize,
+    #[serde(default = "default_rank_order")]
+    pub rank_order: [ParallelAxis; 5],
+    #[serde(default = "default_rank_coordinates")]
+    pub coordinates: RankCoordinates,
+}
+
+fn default_rank_order() -> [ParallelAxis; 5] {
+    ParallelTopology::new(1, 1, 1, 1, 1)
+        .expect("singleton parallel topology is valid")
+        .order()
+}
+
+const fn default_rank_coordinates() -> RankCoordinates {
+    RankCoordinates::ZERO
+}
+
+impl ParallelCheckpointManifest {
+    pub fn new(
+        world_size: usize,
+        global_rank: usize,
+        tensor_model_parallel_size: usize,
+        pipeline_model_parallel_size: usize,
+        data_parallel_size: usize,
+        expert_model_parallel_size: usize,
+        context_parallel_size: usize,
+    ) -> Result<Self> {
+        let topology = ParallelTopology::new(
+            tensor_model_parallel_size,
+            pipeline_model_parallel_size,
+            data_parallel_size,
+            expert_model_parallel_size,
+            context_parallel_size,
+        )?;
+        Self::from_topology(world_size, global_rank, &topology)
+    }
+
+    pub fn from_topology(
+        world_size: usize,
+        global_rank: usize,
+        topology: &ParallelTopology,
+    ) -> Result<Self> {
+        topology.validate_world_size(world_size)?;
+        let coordinates = topology.coordinates(global_rank)?;
+        Ok(Self {
+            world_size,
+            tensor_model_parallel_size: topology.tensor_model_parallel_size(),
+            pipeline_model_parallel_size: topology.pipeline_model_parallel_size(),
+            data_parallel_size: topology.data_parallel_size(),
+            expert_model_parallel_size: topology.expert_model_parallel_size(),
+            context_parallel_size: topology.context_parallel_size(),
+            global_rank,
+            tensor_model_parallel_rank: coordinates.tensor,
+            rank_order: topology.order(),
+            coordinates,
+        })
+    }
+
+    pub fn from_env() -> Result<Self> {
+        let world_size = env_usize(&["WORLD_SIZE"], 1)?;
+        let global_rank = env_usize(&["RANK"], 0)?;
+        let topology = ParallelTopology::from_env_with_world_size(world_size)?;
+        Self::from_topology(world_size, global_rank, &topology)
+    }
+
+    fn is_distributed(&self) -> bool {
+        self.world_size > 1
+    }
+
+    fn legacy_fields_match(&self, other: &Self) -> bool {
+        self.world_size == other.world_size
+            && self.tensor_model_parallel_size == other.tensor_model_parallel_size
+            && self.pipeline_model_parallel_size == other.pipeline_model_parallel_size
+            && self.data_parallel_size == other.data_parallel_size
+            && self.expert_model_parallel_size == other.expert_model_parallel_size
+            && self.context_parallel_size == other.context_parallel_size
+            && self.global_rank == other.global_rank
+            && self.tensor_model_parallel_rank == other.tensor_model_parallel_rank
+    }
+
+    fn replica_identity(&self) -> String {
+        format!(
+            "global-rank-{}-tp-rank-{}",
+            self.global_rank, self.tensor_model_parallel_rank
+        )
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct TensorShardManifest {
+    pub file: String,
+    pub tensor_name: String,
+    pub state: String,
+    #[serde(default)]
+    pub adapter_id: Option<i64>,
+    pub global_lora_rank: i64,
+    pub global_shape: Vec<i64>,
+    pub local_shape: Vec<i64>,
+    pub partition_axis: usize,
+    #[serde(default)]
+    pub layout: LoraTpShardLayout,
+    #[serde(default)]
+    pub replicated: bool,
+    pub global_offset: Vec<i64>,
+    /// Non-contiguous mappings from rank-local storage into the original
+    /// global tensor. Empty for ordinary contiguous shards.
+    #[serde(default)]
+    pub segments: Vec<TensorShardSegmentManifest>,
+    /// V5 supports independent placements on multiple parallel axes.
+    #[serde(default)]
+    pub placements: Vec<TensorShardPlacementManifest>,
+    #[serde(default)]
+    pub replicated_axes: Vec<ParallelAxis>,
+    pub replica_identity: String,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct TensorShardPlacementManifest {
+    pub parallel_axis: ParallelAxis,
+    pub tensor_axis: usize,
+    pub global_size: i64,
+    pub local_size: i64,
+    pub global_offset: i64,
+    #[serde(default)]
+    pub segments: Vec<TensorShardSegmentManifest>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct TensorShardSegmentManifest {
+    pub local_offset: i64,
+    pub global_offset: i64,
+    pub length: i64,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CheckpointManifest {
+    pub format: String,
+    /// New v5 writers publish a compact receipt after this manifest. Readers
+    /// use the receipt set for O(world-size) preflight metadata.
+    #[serde(default)]
+    pub rank_receipt_version: Option<u32>,
+    /// Shared logical generation across every rank in one distributed save.
+    #[serde(default)]
+    pub checkpoint_generation: Option<String>,
+    /// Session/transport progress. Dynamic-only requests advance this clock.
+    pub step: u64,
+    /// Fixed-adapter Adam bias-correction clock. Older checkpoints used
+    /// `step` for both meanings, so a missing value falls back to `step`.
+    #[serde(default)]
+    pub fixed_optimizer_step: Option<u64>,
+    pub loss: f64,
+    pub model_path: String,
+    pub lora_rank: i64,
+    pub lora_alpha: f64,
+    pub files: Vec<String>,
+    /// Stable content digests bind an atomically published manifest to the
+    /// tensor files it describes. Older checkpoints omit this field.
+    #[serde(default)]
+    pub file_digests: BTreeMap<String, String>,
+    #[serde(default)]
+    pub dynamic_adapters: Vec<DynamicAdapterManifest>,
+    #[serde(default)]
+    pub parallel: Option<ParallelCheckpointManifest>,
+    #[serde(default)]
+    pub tensor_shards: Vec<TensorShardManifest>,
+    #[serde(default)]
+    pub fixed_shard_layouts: Vec<LoraTpShardLayout>,
+    #[serde(default)]
+    pub fixed_slot_identities: Vec<LoraSlotIdentity>,
+    /// V6 keeps the immutable global fixed-adapter signature separate from
+    /// the stage-local tensor inventory above.
+    #[serde(default)]
+    pub fixed_target_layers: Vec<usize>,
+    #[serde(default)]
+    pub fixed_target_modules: Vec<String>,
+    #[serde(default)]
+    pub pipeline_stage: Option<PipelineStageCheckpointManifest>,
+}
+
+impl CheckpointManifest {
+    pub fn effective_fixed_optimizer_step(&self) -> u64 {
+        self.fixed_optimizer_step.unwrap_or(self.step)
+    }
+}
+
+fn is_current_distributed_checkpoint(format: &str) -> bool {
+    format == TP_CHECKPOINT_FORMAT || format == STAGE_UNION_CHECKPOINT_FORMAT
+}
+
+fn validate_stage_union_metadata(
+    parallel: &ParallelCheckpointManifest,
+    metadata: &StageUnionCheckpointMetadata,
+    fixed_slot_identities: &[LoraSlotIdentity],
+    dynamic_adapters: &[DynamicAdapterCheckpoint],
+) -> Result<()> {
+    let stage = &metadata.pipeline_stage;
+    if parallel.pipeline_model_parallel_size <= 1 {
+        bail!("stage-union checkpoint requires pipeline_model_parallel_size > 1");
+    }
+    if stage.pipeline_size != parallel.pipeline_model_parallel_size
+        || stage.pipeline_rank != parallel.coordinates.pipeline
+    {
+        bail!(
+            "stage-union checkpoint stage {}/{} does not match topology stage {}/{}",
+            stage.pipeline_rank,
+            stage.pipeline_size,
+            parallel.coordinates.pipeline,
+            parallel.pipeline_model_parallel_size
+        );
+    }
+    if stage.global_num_layers < stage.pipeline_size
+        || stage.layer_start >= stage.layer_end
+        || stage.layer_end > stage.global_num_layers
+    {
+        bail!(
+            "invalid stage-union layer range {}..{} for {} global layers",
+            stage.layer_start,
+            stage.layer_end,
+            stage.global_num_layers
+        );
+    }
+    if metadata
+        .fixed_target_layers
+        .iter()
+        .any(|layer| *layer >= stage.global_num_layers)
+    {
+        bail!("stage-union fixed target layer is outside the global model");
+    }
+    let local_identity_is_valid = |identity: &LoraSlotIdentity| {
+        identity.layer >= stage.layer_start && identity.layer < stage.layer_end
+    };
+    if fixed_slot_identities
+        .iter()
+        .any(|identity| !local_identity_is_valid(identity))
+    {
+        bail!("stage-union fixed slot identity is not owned by this pipeline stage");
+    }
+    for adapter in dynamic_adapters {
+        if adapter
+            .manifest
+            .target_layers
+            .iter()
+            .any(|layer| *layer >= stage.global_num_layers)
+        {
+            bail!(
+                "stage-union dynamic adapter {} target layer is outside the global model",
+                adapter.manifest.id
+            );
+        }
+        if adapter
+            .manifest
+            .slot_identities
+            .iter()
+            .any(|identity| !local_identity_is_valid(identity))
+        {
+            bail!(
+                "stage-union dynamic adapter {} contains a slot not owned by this pipeline stage",
+                adapter.manifest.id
+            );
+        }
+        if adapter.manifest.slot_identities.len() != adapter.manifest.parameter_count
+            || adapter.manifest.shard_layouts.len() != adapter.manifest.parameter_count
+        {
+            bail!(
+                "stage-union dynamic adapter {} local slot metadata count mismatch",
+                adapter.manifest.id
+            );
+        }
+        let unique_slots = adapter
+            .manifest
+            .slot_identities
+            .iter()
+            .map(|identity| (identity.index, identity.layer, identity.module.as_str()))
+            .collect::<std::collections::BTreeSet<_>>();
+        if unique_slots.len() != adapter.manifest.slot_identities.len() {
+            bail!(
+                "stage-union dynamic adapter {} local slot identities must be unique",
+                adapter.manifest.id
+            );
+        }
+    }
+    Ok(())
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+struct CheckpointRankReceipt {
+    format: String,
+    checkpoint_format: String,
+    checkpoint_generation: String,
+    global_rank: usize,
+    parallel: ParallelCheckpointManifest,
+    manifest_digest: String,
+    generation_digest: String,
+    shard_identity_digest: String,
+    shard_count: usize,
+    shard_identities_unique: bool,
+    files: Vec<String>,
+    file_digests: BTreeMap<String, String>,
+    all_files_declared: bool,
+    data_replica_metadata_complete: bool,
+    #[serde(default)]
+    pipeline_stage: Option<PipelineStageCheckpointManifest>,
+    #[serde(default)]
+    stage_state_digest: String,
+}
+
+#[derive(Serialize)]
+struct CheckpointGenerationIdentity<'a> {
+    checkpoint_generation: &'a Option<String>,
+    step: u64,
+    fixed_optimizer_step: u64,
+    model_path: &'a str,
+    lora_rank: i64,
+    lora_alpha_bits: u64,
+    files: &'a [String],
+    dynamic_adapters: &'a [DynamicAdapterManifest],
+    fixed_shard_layouts: &'a [LoraTpShardLayout],
+    fixed_slot_identities: &'a [LoraSlotIdentity],
+}
+
+#[derive(Serialize)]
+struct StageUnionGenerationIdentity<'a> {
+    checkpoint_generation: &'a Option<String>,
+    step: u64,
+    fixed_optimizer_step: u64,
+    model_path: &'a str,
+    lora_rank: i64,
+    lora_alpha_bits: u64,
+    files: &'a [String],
+    fixed_target_layers: &'a [usize],
+    fixed_target_modules: &'a [String],
+    dynamic_adapters: Vec<StageUnionDynamicAdapterIdentity<'a>>,
+}
+
+#[derive(Serialize)]
+struct StageUnionDynamicAdapterIdentity<'a> {
+    id: i64,
+    rank: i64,
+    alpha_bits: u64,
+    optimizer_step: u64,
+    optimizer_lr_bits: Option<u64>,
+    optimizer_beta1_bits: Option<u64>,
+    optimizer_beta2_bits: Option<u64>,
+    optimizer_eps_bits: Option<u64>,
+    target_layers: &'a [usize],
+    target_modules: &'a [String],
+}
+
+/// Stage-union generation identity used before tenant-local beta/epsilon
+/// metadata was added. Keep this wire shape for old rank receipts whose
+/// manifests omit all three optional fields.
+#[derive(Serialize)]
+struct LegacyStageUnionDynamicAdapterIdentity<'a> {
+    id: i64,
+    rank: i64,
+    alpha_bits: u64,
+    optimizer_step: u64,
+    optimizer_lr_bits: Option<u64>,
+    target_layers: &'a [usize],
+    target_modules: &'a [String],
+}
+
+#[derive(Serialize)]
+struct LegacyStageUnionGenerationIdentity<'a> {
+    checkpoint_generation: &'a Option<String>,
+    step: u64,
+    fixed_optimizer_step: u64,
+    model_path: &'a str,
+    lora_rank: i64,
+    lora_alpha_bits: u64,
+    files: &'a [String],
+    fixed_target_layers: &'a [usize],
+    fixed_target_modules: &'a [String],
+    dynamic_adapters: Vec<LegacyStageUnionDynamicAdapterIdentity<'a>>,
+}
+
+#[derive(Serialize)]
+struct StageUnionLocalStateIdentity<'a> {
+    fixed_shard_layouts: &'a [LoraTpShardLayout],
+    fixed_slot_identities: &'a [LoraSlotIdentity],
+    dynamic_adapters: Vec<StageUnionLocalAdapterIdentity<'a>>,
+}
+
+#[derive(Serialize)]
+struct StageUnionLocalAdapterIdentity<'a> {
+    id: i64,
+    shard_layouts: &'a [LoraTpShardLayout],
+    slot_identities: &'a [LoraSlotIdentity],
+    parameter_count: usize,
+    optimizer_count: usize,
+}
+
+pub fn validate_fixed_tp_resume(
+    manifest: &CheckpointManifest,
+    expected_layouts: &[LoraTpShardLayout],
+    expected_identities: &[LoraSlotIdentity],
+) -> Result<()> {
+    if is_legacy_tensor_parallel_checkpoint(manifest) {
+        if expected_layouts
+            .iter()
+            .any(|layout| *layout != LoraTpShardLayout::LatentRank)
+        {
+            bail!(
+                "legacy tensor-parallel v3 checkpoints cannot restore fixed projection-aware attention LoRA; use a v4 checkpoint or migrate the adapter from a merged artifact"
+            );
+        }
+        return Ok(());
+    }
+    if manifest.fixed_shard_layouts != expected_layouts {
+        bail!(
+            "fixed LoRA shard layouts do not match the current runtime slots: checkpoint={:?}, runtime={expected_layouts:?}",
+            manifest.fixed_shard_layouts
+        );
+    }
+    if manifest.fixed_slot_identities != expected_identities {
+        bail!(
+            "fixed LoRA slot identities do not match the current runtime slots: checkpoint={:?}, runtime={expected_identities:?}",
+            manifest.fixed_slot_identities
+        );
+    }
+    Ok(())
+}
+
+pub fn validate_dynamic_tp_resume(
+    manifest: &CheckpointManifest,
+    adapter_id: i64,
+    saved_layouts: &[LoraTpShardLayout],
+    expected_layouts: &[LoraTpShardLayout],
+    expected_identities: &[LoraSlotIdentity],
+) -> Result<()> {
+    if is_legacy_tensor_parallel_checkpoint(manifest) {
+        if expected_layouts
+            .iter()
+            .any(|layout| *layout != LoraTpShardLayout::LatentRank)
+        {
+            bail!(
+                "legacy tensor-parallel v3 checkpoint adapter {adapter_id} contains projection-aware attention LoRA that cannot be restored; use a v4 checkpoint or migrate the adapter from a merged artifact"
+            );
+        }
+        return Ok(());
+    }
+    if saved_layouts != expected_layouts {
+        bail!(
+            "dynamic adapter {adapter_id} shard layouts do not match the current runtime slots: checkpoint={saved_layouts:?}, runtime={expected_layouts:?}"
+        );
+    }
+    if is_current_distributed_checkpoint(&manifest.format) {
+        let saved = manifest
+            .dynamic_adapters
+            .iter()
+            .find(|adapter| adapter.id == adapter_id)
+            .with_context(|| format!("dynamic adapter {adapter_id} is missing from manifest"))?;
+        if saved.slot_identities != expected_identities {
+            bail!(
+                "dynamic adapter {adapter_id} slot identities do not match the current runtime slots: checkpoint={:?}, runtime={expected_identities:?}",
+                saved.slot_identities
+            );
+        }
+    }
+    Ok(())
+}
+
+pub fn fixed_restore_slot_indices(
+    saved_a_count: usize,
+    saved_b_count: usize,
+    active_slot_indices: &[usize],
+    native_slot_count: usize,
+) -> Result<Vec<usize>> {
+    if saved_a_count != saved_b_count {
+        bail!("checkpoint fixed LoRA A/B count mismatch: {saved_a_count}/{saved_b_count}");
+    }
+    if saved_a_count == active_slot_indices.len() {
+        return Ok(active_slot_indices.to_vec());
+    }
+    if saved_a_count == native_slot_count {
+        // v1/v2 checkpoints stored inactive positional placeholders.
+        return Ok((0..native_slot_count).collect());
+    }
+    bail!(
+        "checkpoint LoRA slot count mismatch: checkpoint A/B={saved_a_count}/{saved_b_count}, active={}, native={native_slot_count}",
+        active_slot_indices.len()
+    )
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct DynamicAdapterManifest {
+    pub id: i64,
+    pub rank: i64,
+    pub alpha: f64,
+    /// Adam bias-correction clock for this tenant. It is independent from
+    /// the session-wide `CheckpointManifest::step` and other adapters.
+    /// Missing values in older v2 manifests default to zero.
+    #[serde(default)]
+    pub optimizer_step: u64,
+    /// Effective Adam learning rate for this tenant. Older manifests omit it
+    /// and restore with the training context learning rate.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub optimizer_lr: Option<f64>,
+    /// Tenant-local Adam beta1. Older manifests inherit the session value.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub optimizer_beta1: Option<f64>,
+    /// Tenant-local Adam beta2. Older manifests inherit the session value.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub optimizer_beta2: Option<f64>,
+    /// Tenant-local Adam epsilon. Older manifests inherit the session value.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub optimizer_eps: Option<f64>,
+    pub target_layers: Vec<usize>,
+    pub target_modules: Vec<String>,
+    #[serde(default)]
+    pub shard_layouts: Vec<LoraTpShardLayout>,
+    #[serde(default)]
+    pub slot_identities: Vec<LoraSlotIdentity>,
+    pub parameter_count: usize,
+    pub optimizer_count: usize,
+}
+
+pub struct DynamicAdapterCheckpoint {
+    pub manifest: DynamicAdapterManifest,
+    pub lora_a: Vec<Tensor>,
+    pub lora_b: Vec<Tensor>,
+    pub adam_m: Vec<Tensor>,
+    pub adam_v: Vec<Tensor>,
+}
+
+pub struct CheckpointData {
+    pub manifest: CheckpointManifest,
+    pub lora_a: Vec<Tensor>,
+    pub lora_b: Vec<Tensor>,
+    pub adam_m: Vec<Tensor>,
+    pub adam_v: Vec<Tensor>,
+    pub dynamic_adapters: Vec<DynamicAdapterCheckpoint>,
+}
+
+pub struct MergedLoraSlot {
+    pub identity: LoraSlotIdentity,
+    pub lora_a: Tensor,
+    pub lora_b: Tensor,
+}
+
+pub struct MergedAdapterCheckpoint {
+    pub model_path: String,
+    pub step: u64,
+    pub rank: i64,
+    pub alpha: f64,
+    pub optimizer_step: u64,
+    pub target_layers: Vec<usize>,
+    pub target_modules: Vec<String>,
+    pub slots: Vec<MergedLoraSlot>,
+}
+
+struct TensorFragment {
+    global_ranges: Vec<(i64, i64)>,
+    tensor: Tensor,
+}
+
+/// Merge one fixed (`None`) or dynamic adapter from a complete v5 rank set.
+/// This is an offline artifact operation, not part of the training hot path.
+pub fn merge_distributed_adapter_checkpoint(
+    root: &Path,
+    adapter_id: Option<i64>,
+) -> Result<MergedAdapterCheckpoint> {
+    let rank_zero_path = rank_checkpoint_dir(root, 0).join("manifest.json");
+    let rank_zero: CheckpointManifest = serde_json::from_str(
+        &std::fs::read_to_string(&rank_zero_path)
+            .with_context(|| format!("read {}", rank_zero_path.display()))?,
+    )
+    .with_context(|| format!("parse {}", rank_zero_path.display()))?;
+    if rank_zero.format != TP_CHECKPOINT_FORMAT {
+        bail!("distributed adapter merge requires a v5 checkpoint");
+    }
+    let topology = rank_zero
+        .parallel
+        .as_ref()
+        .context("rank 0 checkpoint is missing parallel topology")?;
+    let world_size = topology.world_size;
+    let rank_order = topology
+        .rank_order
+        .iter()
+        .map(|axis| axis.name())
+        .collect::<Vec<_>>()
+        .join("-");
+    let expected_topology = ParallelTopology::with_order(
+        topology.tensor_model_parallel_size,
+        topology.pipeline_model_parallel_size,
+        topology.data_parallel_size,
+        topology.expert_model_parallel_size,
+        topology.context_parallel_size,
+        &rank_order,
+    )?;
+    expected_topology.validate_world_size(world_size)?;
+    let (rank, alpha, optimizer_step, target_layers, target_modules, identities) = match adapter_id
+    {
+        None | Some(0) => (
+            rank_zero.lora_rank,
+            rank_zero.lora_alpha,
+            rank_zero.effective_fixed_optimizer_step(),
+            rank_zero
+                .fixed_slot_identities
+                .iter()
+                .map(|identity| identity.layer)
+                .collect::<std::collections::BTreeSet<_>>()
+                .into_iter()
+                .collect(),
+            rank_zero
+                .fixed_slot_identities
+                .iter()
+                .map(|identity| identity.module.clone())
+                .collect::<std::collections::BTreeSet<_>>()
+                .into_iter()
+                .collect(),
+            rank_zero.fixed_slot_identities.clone(),
+        ),
+        Some(id) => {
+            let adapter = rank_zero
+                .dynamic_adapters
+                .iter()
+                .find(|adapter| adapter.id == id)
+                .with_context(|| format!("rank 0 checkpoint has no dynamic adapter {id}"))?;
+            (
+                adapter.rank,
+                adapter.alpha,
+                adapter.optimizer_step,
+                adapter.target_layers.clone(),
+                adapter.target_modules.clone(),
+                adapter.slot_identities.clone(),
+            )
+        }
+    };
+    if identities.is_empty() {
+        bail!("distributed adapter merge requires explicit slot identities");
+    }
+
+    let selected_id = adapter_id.filter(|id| *id != 0);
+    let mut fragments: BTreeMap<(String, String), (Vec<i64>, Vec<TensorFragment>)> =
+        BTreeMap::new();
+    let mut seen_ranks = std::collections::BTreeSet::new();
+    for global_rank in 0..world_size {
+        let rank_dir = rank_checkpoint_dir(root, global_rank);
+        let manifest_path = rank_dir.join("manifest.json");
+        let manifest: CheckpointManifest = serde_json::from_str(
+            &std::fs::read_to_string(&manifest_path)
+                .with_context(|| format!("read {}", manifest_path.display()))?,
+        )
+        .with_context(|| format!("parse {}", manifest_path.display()))?;
+        validate_manifest_file_digests(&rank_dir, &manifest)?;
+        let parallel = manifest
+            .parallel
+            .as_ref()
+            .with_context(|| format!("rank {global_rank} checkpoint is missing topology"))?;
+        if manifest.format != TP_CHECKPOINT_FORMAT
+            || parallel.world_size != world_size
+            || parallel.global_rank != global_rank
+            || parallel.rank_order != topology.rank_order
+            || parallel.tensor_model_parallel_size != topology.tensor_model_parallel_size
+            || parallel.pipeline_model_parallel_size != topology.pipeline_model_parallel_size
+            || parallel.data_parallel_size != topology.data_parallel_size
+            || parallel.expert_model_parallel_size != topology.expert_model_parallel_size
+            || parallel.context_parallel_size != topology.context_parallel_size
+            || parallel.coordinates != expected_topology.coordinates(global_rank)?
+            || manifest.step != rank_zero.step
+            || manifest.effective_fixed_optimizer_step()
+                != rank_zero.effective_fixed_optimizer_step()
+            || manifest.model_path != rank_zero.model_path
+        {
+            bail!("rank {global_rank} checkpoint metadata is inconsistent with rank 0");
+        }
+        validate_checkpoint_generation(&rank_zero, &manifest, global_rank)?;
+        if !seen_ranks.insert((
+            parallel.global_rank,
+            parallel.coordinates.tensor,
+            parallel.coordinates.pipeline,
+            parallel.coordinates.data,
+            parallel.coordinates.expert,
+            parallel.coordinates.context,
+        )) {
+            bail!("duplicate distributed checkpoint coordinates at rank {global_rank}");
+        }
+        match selected_id {
+            None => {
+                if manifest.lora_rank != rank_zero.lora_rank
+                    || manifest.lora_alpha != rank_zero.lora_alpha
+                    || manifest.fixed_shard_layouts != rank_zero.fixed_shard_layouts
+                    || manifest.fixed_slot_identities != rank_zero.fixed_slot_identities
+                {
+                    bail!("rank {global_rank} fixed adapter metadata differs from rank 0");
+                }
+            }
+            Some(id) => {
+                let expected_adapter = rank_zero
+                    .dynamic_adapters
+                    .iter()
+                    .find(|adapter| adapter.id == id)
+                    .expect("selected dynamic adapter was validated on rank 0");
+                let adapter = manifest
+                    .dynamic_adapters
+                    .iter()
+                    .find(|adapter| adapter.id == id)
+                    .with_context(|| format!("rank {global_rank} has no dynamic adapter {id}"))?;
+                if adapter != expected_adapter {
+                    bail!("rank {global_rank} dynamic adapter {id} metadata differs from rank 0");
+                }
+            }
+        }
+        let tensors = read_named_tensors(&rank_dir.join("adapter.safetensors"))?;
+        for shard in manifest.tensor_shards.iter().filter(|shard| {
+            shard.file == "adapter.safetensors"
+                && (shard.state == "lora_a" || shard.state == "lora_b")
+                && shard.adapter_id == selected_id
+        }) {
+            let tensor = tensors.get(&shard.tensor_name).with_context(|| {
+                format!(
+                    "rank {global_rank} adapter is missing tensor {}",
+                    shard.tensor_name
+                )
+            })?;
+            if tensor.size() != shard.local_shape {
+                bail!(
+                    "rank {global_rank} tensor {} shape {:?} does not match manifest {:?}",
+                    shard.tensor_name,
+                    tensor.size(),
+                    shard.local_shape
+                );
+            }
+            let key = (shard.state.clone(), shard.tensor_name.clone());
+            let entry = fragments
+                .entry(key)
+                .or_insert_with(|| (shard.global_shape.clone(), Vec::new()));
+            if entry.0 != shard.global_shape {
+                bail!("global shape mismatch for tensor {}", shard.tensor_name);
+            }
+            entry.1.extend(tensor_fragments(shard, tensor)?);
+        }
+    }
+    if seen_ranks.len() != world_size {
+        bail!("distributed checkpoint rank set is incomplete");
+    }
+
+    let merged = fragments
+        .into_iter()
+        .map(|(key, (shape, fragments))| Ok((key, merge_tensor_fragments(&shape, fragments)?)))
+        .collect::<Result<BTreeMap<_, _>>>()?;
+    let prefix = selected_id
+        .map(|id| format!("dynamic_{id}_"))
+        .unwrap_or_default();
+    let mut slots = Vec::with_capacity(identities.len());
+    for (index, identity) in identities.into_iter().enumerate() {
+        let a_name = format!("{prefix}a_{index}");
+        let b_name = format!("{prefix}b_{index}");
+        let lora_a = merged
+            .get(&("lora_a".to_string(), a_name.clone()))
+            .with_context(|| format!("merged adapter is missing {a_name}"))?
+            .shallow_clone();
+        let lora_b = merged
+            .get(&("lora_b".to_string(), b_name.clone()))
+            .with_context(|| format!("merged adapter is missing {b_name}"))?
+            .shallow_clone();
+        slots.push(MergedLoraSlot {
+            identity,
+            lora_a,
+            lora_b,
+        });
+    }
+    Ok(MergedAdapterCheckpoint {
+        model_path: rank_zero.model_path,
+        step: rank_zero.step,
+        rank,
+        alpha,
+        optimizer_step,
+        target_layers,
+        target_modules,
+        slots,
+    })
+}
+
+fn tensor_fragments(shard: &TensorShardManifest, tensor: &Tensor) -> Result<Vec<TensorFragment>> {
+    let dimensions = shard.local_shape.len();
+    let mut mappings = shard
+        .local_shape
+        .iter()
+        .map(|size| vec![(0, 0, *size)])
+        .collect::<Vec<_>>();
+    let mut placed_axes = std::collections::BTreeSet::new();
+    for placement in &shard.placements {
+        if placement.tensor_axis >= dimensions || !placed_axes.insert(placement.tensor_axis) {
+            bail!(
+                "invalid or duplicate placement axis for {}",
+                shard.tensor_name
+            );
+        }
+        mappings[placement.tensor_axis] = if placement.segments.is_empty() {
+            vec![(0, placement.global_offset, placement.local_size)]
+        } else {
+            placement
+                .segments
+                .iter()
+                .map(|segment| (segment.local_offset, segment.global_offset, segment.length))
+                .collect()
+        };
+    }
+    let mut selections = Vec::new();
+    cartesian_axis_mappings(&mappings, 0, &mut Vec::new(), &mut selections);
+    let mut fragments = Vec::with_capacity(selections.len());
+    for selection in selections {
+        let mut local = tensor.shallow_clone();
+        for (axis, (local_offset, _, length)) in selection.iter().copied().enumerate() {
+            if local_offset < 0 || length <= 0 || local_offset + length > shard.local_shape[axis] {
+                bail!("local placement is out of bounds for {}", shard.tensor_name);
+            }
+            local = local.narrow(axis as i64, local_offset, length);
+        }
+        let global_ranges = selection
+            .iter()
+            .enumerate()
+            .map(|(axis, (_, global_offset, length))| {
+                if *global_offset < 0
+                    || *length <= 0
+                    || *global_offset + *length > shard.global_shape[axis]
+                {
+                    bail!(
+                        "global placement is out of bounds for {}",
+                        shard.tensor_name
+                    );
+                }
+                Ok((*global_offset, *length))
+            })
+            .collect::<Result<Vec<_>>>()?;
+        fragments.push(TensorFragment {
+            global_ranges,
+            tensor: local,
+        });
+    }
+    Ok(fragments)
+}
+
+fn cartesian_axis_mappings(
+    mappings: &[Vec<(i64, i64, i64)>],
+    axis: usize,
+    current: &mut Vec<(i64, i64, i64)>,
+    output: &mut Vec<Vec<(i64, i64, i64)>>,
+) {
+    if axis == mappings.len() {
+        output.push(current.clone());
+        return;
+    }
+    for mapping in &mappings[axis] {
+        current.push(*mapping);
+        cartesian_axis_mappings(mappings, axis + 1, current, output);
+        current.pop();
+    }
+}
+
+fn merge_tensor_fragments(shape: &[i64], fragments: Vec<TensorFragment>) -> Result<Tensor> {
+    let mut unique: BTreeMap<Vec<(i64, i64)>, Tensor> = BTreeMap::new();
+    for fragment in fragments {
+        if let Some(existing) = unique.get(&fragment.global_ranges) {
+            if !existing.allclose(&fragment.tensor, 0.0, 0.0, false) {
+                bail!("replicated distributed checkpoint fragments differ");
+            }
+        } else {
+            unique.insert(fragment.global_ranges, fragment.tensor);
+        }
+    }
+    let regions = unique.keys().collect::<Vec<_>>();
+    for left in 0..regions.len() {
+        for right in (left + 1)..regions.len() {
+            let overlaps = regions[left].iter().zip(regions[right].iter()).all(
+                |(&(left_start, left_len), &(right_start, right_len))| {
+                    left_start < right_start + right_len && right_start < left_start + left_len
+                },
+            );
+            if overlaps {
+                bail!("distributed checkpoint fragments overlap");
+            }
+        }
+    }
+    let covered = unique
+        .keys()
+        .map(|ranges| ranges.iter().map(|(_, length)| *length).product::<i64>())
+        .sum::<i64>();
+    let expected = shape.iter().product::<i64>();
+    if covered != expected {
+        bail!("distributed checkpoint fragments cover {covered} elements, expected {expected}");
+    }
+    let output = Tensor::zeros(shape, (tch::Kind::Float, tch::Device::Cpu));
+    for (ranges, source) in unique {
+        let mut target = output.shallow_clone();
+        for (axis, (offset, length)) in ranges.into_iter().enumerate() {
+            target = target.narrow(axis as i64, offset, length);
+        }
+        target.copy_(&source.to_kind(tch::Kind::Float));
+    }
+    Ok(output)
+}
+
+/// Save checkpoint to a directory.
+/// Creates: manifest.json, adapter.safetensors, optimizer.safetensors
+pub fn save_checkpoint(
+    dir: &Path,
+    step: u64,
+    loss: f64,
+    model_path: &str,
+    lora_rank: i64,
+    lora_alpha: f64,
+    lora_a: &[Tensor],
+    lora_b: &[Tensor],
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+) -> Result<()> {
+    save_checkpoint_with_dynamic(
+        dir,
+        step,
+        loss,
+        model_path,
+        lora_rank,
+        lora_alpha,
+        lora_a,
+        lora_b,
+        adam_m,
+        adam_v,
+        &[],
+    )
+}
+
+pub fn save_checkpoint_with_dynamic(
+    dir: &Path,
+    step: u64,
+    loss: f64,
+    model_path: &str,
+    lora_rank: i64,
+    lora_alpha: f64,
+    lora_a: &[Tensor],
+    lora_b: &[Tensor],
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+    dynamic_adapters: &[DynamicAdapterCheckpoint],
+) -> Result<()> {
+    save_checkpoint_with_dynamic_at(
+        dir,
+        step,
+        if lora_a.is_empty() { 0 } else { step },
+        loss,
+        model_path,
+        lora_rank,
+        lora_alpha,
+        lora_a,
+        lora_b,
+        adam_m,
+        adam_v,
+        dynamic_adapters,
+        &[],
+        &[],
+        None,
+        None,
+        None,
+    )
+}
+
+#[allow(clippy::too_many_arguments)]
+pub fn save_checkpoint_with_dynamic_for_topology(
+    dir: &Path,
+    step: u64,
+    loss: f64,
+    model_path: &str,
+    lora_rank: i64,
+    lora_alpha: f64,
+    lora_a: &[Tensor],
+    lora_b: &[Tensor],
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+    dynamic_adapters: &[DynamicAdapterCheckpoint],
+    fixed_shard_layouts: &[LoraTpShardLayout],
+    fixed_slot_identities: &[LoraSlotIdentity],
+    parallel: &ParallelCheckpointManifest,
+) -> Result<()> {
+    save_checkpoint_with_dynamic_and_fixed_step_for_topology(
+        dir,
+        step,
+        if lora_a.is_empty() { 0 } else { step },
+        loss,
+        model_path,
+        lora_rank,
+        lora_alpha,
+        lora_a,
+        lora_b,
+        adam_m,
+        adam_v,
+        dynamic_adapters,
+        fixed_shard_layouts,
+        fixed_slot_identities,
+        parallel,
+    )
+}
+
+#[allow(clippy::too_many_arguments)]
+pub fn save_checkpoint_with_dynamic_and_fixed_step_for_topology(
+    dir: &Path,
+    step: u64,
+    fixed_optimizer_step: u64,
+    loss: f64,
+    model_path: &str,
+    lora_rank: i64,
+    lora_alpha: f64,
+    lora_a: &[Tensor],
+    lora_b: &[Tensor],
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+    dynamic_adapters: &[DynamicAdapterCheckpoint],
+    fixed_shard_layouts: &[LoraTpShardLayout],
+    fixed_slot_identities: &[LoraSlotIdentity],
+    parallel: &ParallelCheckpointManifest,
+) -> Result<()> {
+    if parallel.is_distributed() {
+        bail!(
+            "distributed checkpoint save requires an explicit unique generation from its coordinator"
+        );
+    }
+    save_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+        dir,
+        step,
+        fixed_optimizer_step,
+        loss,
+        model_path,
+        lora_rank,
+        lora_alpha,
+        lora_a,
+        lora_b,
+        adam_m,
+        adam_v,
+        dynamic_adapters,
+        fixed_shard_layouts,
+        fixed_slot_identities,
+        parallel,
+        None,
+    )
+}
+
+#[allow(clippy::too_many_arguments)]
+pub fn save_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+    dir: &Path,
+    step: u64,
+    fixed_optimizer_step: u64,
+    loss: f64,
+    model_path: &str,
+    lora_rank: i64,
+    lora_alpha: f64,
+    lora_a: &[Tensor],
+    lora_b: &[Tensor],
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+    dynamic_adapters: &[DynamicAdapterCheckpoint],
+    fixed_shard_layouts: &[LoraTpShardLayout],
+    fixed_slot_identities: &[LoraSlotIdentity],
+    parallel: &ParallelCheckpointManifest,
+    checkpoint_generation: Option<&str>,
+) -> Result<()> {
+    if !parallel.is_distributed() {
+        return save_checkpoint_with_dynamic_at(
+            dir,
+            step,
+            fixed_optimizer_step,
+            loss,
+            model_path,
+            lora_rank,
+            lora_alpha,
+            lora_a,
+            lora_b,
+            adam_m,
+            adam_v,
+            dynamic_adapters,
+            &[],
+            &[],
+            None,
+            None,
+            None,
+        );
+    }
+    let checkpoint_generation = checkpoint_generation
+        .filter(|generation| !generation.is_empty())
+        .context("distributed checkpoint generation must be non-empty")?;
+    validate_checkpoint_generation_value(checkpoint_generation)?;
+    let rank_dir = rank_checkpoint_dir(dir, parallel.global_rank);
+    let manifest_path = rank_dir.join("manifest.json");
+    if manifest_path.is_file() {
+        let previous = read_checkpoint_manifest(&manifest_path)?;
+        if previous.checkpoint_generation.as_deref() == Some(checkpoint_generation) {
+            bail!(
+                "distributed checkpoint generation {checkpoint_generation} has already been used for rank {}",
+                parallel.global_rank
+            );
+        }
+    }
+    save_checkpoint_with_dynamic_at(
+        &rank_dir,
+        step,
+        fixed_optimizer_step,
+        loss,
+        model_path,
+        lora_rank,
+        lora_alpha,
+        lora_a,
+        lora_b,
+        adam_m,
+        adam_v,
+        dynamic_adapters,
+        fixed_shard_layouts,
+        fixed_slot_identities,
+        Some(parallel),
+        Some(checkpoint_generation),
+        None,
+    )
+}
+
+#[allow(clippy::too_many_arguments)]
+pub fn save_stage_union_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+    dir: &Path,
+    step: u64,
+    fixed_optimizer_step: u64,
+    loss: f64,
+    model_path: &str,
+    lora_rank: i64,
+    lora_alpha: f64,
+    lora_a: &[Tensor],
+    lora_b: &[Tensor],
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+    dynamic_adapters: &[DynamicAdapterCheckpoint],
+    fixed_shard_layouts: &[LoraTpShardLayout],
+    fixed_slot_identities: &[LoraSlotIdentity],
+    parallel: &ParallelCheckpointManifest,
+    checkpoint_generation: &str,
+    stage_union: &StageUnionCheckpointMetadata,
+) -> Result<()> {
+    validate_stage_union_metadata(
+        parallel,
+        stage_union,
+        fixed_slot_identities,
+        dynamic_adapters,
+    )?;
+    let checkpoint_generation = (!checkpoint_generation.is_empty())
+        .then_some(checkpoint_generation)
+        .context("distributed checkpoint generation must be non-empty")?;
+    validate_checkpoint_generation_value(checkpoint_generation)?;
+    let rank_dir = rank_checkpoint_dir(dir, parallel.global_rank);
+    let manifest_path = rank_dir.join("manifest.json");
+    if manifest_path.is_file() {
+        let previous = read_checkpoint_manifest(&manifest_path)?;
+        if previous.checkpoint_generation.as_deref() == Some(checkpoint_generation) {
+            bail!(
+                "distributed checkpoint generation {checkpoint_generation} has already been used for rank {}",
+                parallel.global_rank
+            );
+        }
+    }
+    save_checkpoint_with_dynamic_at(
+        &rank_dir,
+        step,
+        fixed_optimizer_step,
+        loss,
+        model_path,
+        lora_rank,
+        lora_alpha,
+        lora_a,
+        lora_b,
+        adam_m,
+        adam_v,
+        dynamic_adapters,
+        fixed_shard_layouts,
+        fixed_slot_identities,
+        Some(parallel),
+        Some(checkpoint_generation),
+        Some(stage_union),
+    )
+}
+
+fn validate_checkpoint_generation_value(generation: &str) -> Result<()> {
+    if generation.len() > 256
+        || !generation
+            .bytes()
+            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
+    {
+        bail!("distributed checkpoint generation must be at most 256 path-safe ASCII characters");
+    }
+    Ok(())
+}
+
+#[allow(clippy::too_many_arguments)]
+fn save_checkpoint_with_dynamic_at(
+    dir: &Path,
+    step: u64,
+    fixed_optimizer_step: u64,
+    loss: f64,
+    model_path: &str,
+    lora_rank: i64,
+    lora_alpha: f64,
+    lora_a: &[Tensor],
+    lora_b: &[Tensor],
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+    dynamic_adapters: &[DynamicAdapterCheckpoint],
+    fixed_shard_layouts: &[LoraTpShardLayout],
+    fixed_slot_identities: &[LoraSlotIdentity],
+    parallel: Option<&ParallelCheckpointManifest>,
+    checkpoint_generation: Option<&str>,
+    stage_union: Option<&StageUnionCheckpointMetadata>,
+) -> Result<()> {
+    validate_tensor_counts(
+        lora_a,
+        lora_b,
+        adam_m,
+        adam_v,
+        dynamic_adapters,
+        parallel.is_some(),
+    )?;
+    if stage_union.is_some()
+        != parallel.is_some_and(|parallel| parallel.pipeline_model_parallel_size > 1)
+    {
+        bail!("stage-union metadata is required exactly when pipeline parallelism is enabled");
+    }
+    if fixed_optimizer_step > 0
+        && (stage_union.is_none() || !lora_a.is_empty())
+        && (adam_m.is_empty() || adam_v.is_empty())
+    {
+        bail!("fixed optimizer step {fixed_optimizer_step} requires Adam state");
+    }
+    for adapter in dynamic_adapters {
+        if adapter.manifest.optimizer_step > 0
+            && (stage_union.is_none() || adapter.manifest.parameter_count > 0)
+            && (adapter.adam_m.is_empty() || adapter.adam_v.is_empty())
+        {
+            bail!(
+                "dynamic adapter {} optimizer step {} requires Adam state",
+                adapter.manifest.id,
+                adapter.manifest.optimizer_step
+            );
+        }
+    }
+    if parallel.is_some() && fixed_slot_identities.len() != lora_a.len() {
+        bail!(
+            "fixed LoRA slot identity count {} does not match parameter count {}",
+            fixed_slot_identities.len(),
+            lora_a.len()
+        );
+    }
+    let unique_fixed_slots = fixed_slot_identities
+        .iter()
+        .map(|identity| (identity.index, identity.layer, identity.module.as_str()))
+        .collect::<std::collections::BTreeSet<_>>();
+    if unique_fixed_slots.len() != fixed_slot_identities.len() {
+        bail!("fixed LoRA slot identities must be unique");
+    }
+    let tensor_shards = match parallel {
+        Some(parallel) => build_tensor_shard_manifest(
+            parallel,
+            lora_rank,
+            lora_a,
+            lora_b,
+            adam_m,
+            adam_v,
+            dynamic_adapters,
+            fixed_shard_layouts,
+        )?,
+        None => Vec::new(),
+    };
+    std::fs::create_dir_all(dir)
+        .with_context(|| format!("create checkpoint dir {}", dir.display()))?;
+
+    // Save adapter (LoRA A/B) as safetensors
+    let adapter_path = dir.join("adapter.safetensors");
+    let mut adapter_tensors = named_tensors(lora_a, lora_b, "a_", "b_");
+    let mut dynamic_manifests = Vec::with_capacity(dynamic_adapters.len());
+    for adapter in dynamic_adapters {
+        let id = adapter.manifest.id;
+        adapter_tensors.extend(named_tensors(
+            &adapter.lora_a,
+            &adapter.lora_b,
+            &format!("dynamic_{id}_a_"),
+            &format!("dynamic_{id}_b_"),
+        ));
+        dynamic_manifests.push(adapter.manifest.clone());
+    }
+    save_named_tensors(&adapter_path, adapter_tensors)?;
+
+    // Save optimizer state (Adam m/v) as safetensors
+    let optimizer_path = dir.join("optimizer.safetensors");
+    let mut optimizer_tensors = named_tensors(adam_m, adam_v, "a_", "b_");
+    for adapter in dynamic_adapters {
+        let id = adapter.manifest.id;
+        optimizer_tensors.extend(named_tensors(
+            &adapter.adam_m,
+            &adapter.adam_v,
+            &format!("dynamic_{id}_a_"),
+            &format!("dynamic_{id}_b_"),
+        ));
+    }
+    save_named_tensors(&optimizer_path, optimizer_tensors)?;
+
+    let file_digests = [
+        (
+            "adapter.safetensors".to_string(),
+            stable_file_digest(&adapter_path)?,
+        ),
+        (
+            "optimizer.safetensors".to_string(),
+            stable_file_digest(&optimizer_path)?,
+        ),
+    ]
+    .into_iter()
+    .collect();
+
+    // Write manifest
+    let manifest = CheckpointManifest {
+        format: if stage_union.is_some() {
+            STAGE_UNION_CHECKPOINT_FORMAT.to_string()
+        } else if parallel.is_some() {
+            TP_CHECKPOINT_FORMAT.to_string()
+        } else if dynamic_manifests.is_empty() {
+            "rustrain-checkpoint-v1".to_string()
+        } else {
+            "rustrain-checkpoint-v2".to_string()
+        },
+        rank_receipt_version: parallel.map(|_| RANK_RECEIPT_VERSION),
+        checkpoint_generation: checkpoint_generation.map(str::to_string),
+        step,
+        fixed_optimizer_step: Some(fixed_optimizer_step),
+        loss,
+        model_path: model_path.to_string(),
+        lora_rank,
+        lora_alpha,
+        files: vec!["adapter.safetensors".into(), "optimizer.safetensors".into()],
+        file_digests,
+        dynamic_adapters: dynamic_manifests,
+        parallel: parallel.cloned(),
+        tensor_shards,
+        fixed_shard_layouts: fixed_shard_layouts.to_vec(),
+        fixed_slot_identities: fixed_slot_identities.to_vec(),
+        fixed_target_layers: stage_union
+            .map(|metadata| metadata.fixed_target_layers.clone())
+            .unwrap_or_default(),
+        fixed_target_modules: stage_union
+            .map(|metadata| metadata.fixed_target_modules.clone())
+            .unwrap_or_default(),
+        pipeline_stage: stage_union.map(|metadata| metadata.pipeline_stage.clone()),
+    };
+    let manifest_path = dir.join("manifest.json");
+    let manifest_contents = serde_json::to_vec_pretty(&manifest)?;
+    write_atomic(&manifest_path, &manifest_contents).with_context(|| "write manifest.json")?;
+    if parallel.is_some() {
+        let receipt = checkpoint_rank_receipt(&manifest, &manifest_contents)?;
+        write_atomic(
+            &dir.join(RANK_RECEIPT_FILE),
+            &serde_json::to_vec_pretty(&receipt)?,
+        )
+        .with_context(|| format!("write {RANK_RECEIPT_FILE}"))?;
+    }
+
+    tracing::info!(
+        step,
+        loss,
+        path = dir.display().to_string(),
+        "checkpoint saved"
+    );
+    Ok(())
+}
+
+pub fn write_atomic(path: &Path, contents: &[u8]) -> Result<()> {
+    let file_name = path
+        .file_name()
+        .and_then(|name| name.to_str())
+        .context("atomic output path has no UTF-8 file name")?;
+    let partial_path = path.with_file_name(format!(".{file_name}.partial"));
+    std::fs::write(&partial_path, contents)
+        .with_context(|| format!("write partial file {}", partial_path.display()))?;
+    std::fs::rename(&partial_path, path).with_context(|| {
+        format!(
+            "publish atomic file {} as {}",
+            partial_path.display(),
+            path.display()
+        )
+    })?;
+    Ok(())
+}
+
+fn stable_file_digest(path: &Path) -> Result<String> {
+    let mut file = std::fs::File::open(path)
+        .with_context(|| format!("open checkpoint content {}", path.display()))?;
+    let mut first = 0xcbf29ce484222325_u64;
+    let mut second = 0x84222325cbf29ce4_u64;
+    let mut buffer = [0_u8; 64 * 1024];
+    loop {
+        let count = file
+            .read(&mut buffer)
+            .with_context(|| format!("hash checkpoint content {}", path.display()))?;
+        if count == 0 {
+            break;
+        }
+        for &byte in &buffer[..count] {
+            first ^= u64::from(byte);
+            first = first.wrapping_mul(0x100000001b3);
+            second ^= u64::from(byte).wrapping_add(0x9d);
+            second = second.wrapping_mul(0x100000001b3 ^ 0x517cc1b727220a95);
+        }
+    }
+    Ok(format!("fnv128-v1:{first:016x}{second:016x}"))
+}
+
+fn stable_bytes_digest(contents: &[u8]) -> String {
+    let mut first = 0xcbf29ce484222325_u64;
+    let mut second = 0x84222325cbf29ce4_u64;
+    for &byte in contents {
+        first ^= u64::from(byte);
+        first = first.wrapping_mul(0x100000001b3);
+        second ^= u64::from(byte).wrapping_add(0x9d);
+        second = second.wrapping_mul(0x100000001b3 ^ 0x517cc1b727220a95);
+    }
+    format!("fnv128-v1:{first:016x}{second:016x}")
+}
+
+fn checkpoint_generation_digest(manifest: &CheckpointManifest) -> Result<String> {
+    if manifest.format == STAGE_UNION_CHECKPOINT_FORMAT {
+        let use_legacy_identity = manifest.dynamic_adapters.iter().all(|adapter| {
+            adapter.optimizer_beta1.is_none()
+                && adapter.optimizer_beta2.is_none()
+                && adapter.optimizer_eps.is_none()
+        });
+        if use_legacy_identity {
+            let dynamic_adapters = manifest
+                .dynamic_adapters
+                .iter()
+                .map(|adapter| LegacyStageUnionDynamicAdapterIdentity {
+                    id: adapter.id,
+                    rank: adapter.rank,
+                    alpha_bits: adapter.alpha.to_bits(),
+                    optimizer_step: adapter.optimizer_step,
+                    optimizer_lr_bits: adapter.optimizer_lr.map(f64::to_bits),
+                    target_layers: &adapter.target_layers,
+                    target_modules: &adapter.target_modules,
+                })
+                .collect();
+            return Ok(stable_bytes_digest(&serde_json::to_vec(
+                &LegacyStageUnionGenerationIdentity {
+                    checkpoint_generation: &manifest.checkpoint_generation,
+                    step: manifest.step,
+                    fixed_optimizer_step: manifest.effective_fixed_optimizer_step(),
+                    model_path: &manifest.model_path,
+                    lora_rank: manifest.lora_rank,
+                    lora_alpha_bits: manifest.lora_alpha.to_bits(),
+                    files: &manifest.files,
+                    fixed_target_layers: &manifest.fixed_target_layers,
+                    fixed_target_modules: &manifest.fixed_target_modules,
+                    dynamic_adapters,
+                },
+            )?));
+        }
+        let dynamic_adapters = manifest
+            .dynamic_adapters
+            .iter()
+            .map(|adapter| StageUnionDynamicAdapterIdentity {
+                id: adapter.id,
+                rank: adapter.rank,
+                alpha_bits: adapter.alpha.to_bits(),
+                optimizer_step: adapter.optimizer_step,
+                optimizer_lr_bits: adapter.optimizer_lr.map(f64::to_bits),
+                optimizer_beta1_bits: adapter.optimizer_beta1.map(f64::to_bits),
+                optimizer_beta2_bits: adapter.optimizer_beta2.map(f64::to_bits),
+                optimizer_eps_bits: adapter.optimizer_eps.map(f64::to_bits),
+                target_layers: &adapter.target_layers,
+                target_modules: &adapter.target_modules,
+            })
+            .collect();
+        return Ok(stable_bytes_digest(&serde_json::to_vec(
+            &StageUnionGenerationIdentity {
+                checkpoint_generation: &manifest.checkpoint_generation,
+                step: manifest.step,
+                fixed_optimizer_step: manifest.effective_fixed_optimizer_step(),
+                model_path: &manifest.model_path,
+                lora_rank: manifest.lora_rank,
+                lora_alpha_bits: manifest.lora_alpha.to_bits(),
+                files: &manifest.files,
+                fixed_target_layers: &manifest.fixed_target_layers,
+                fixed_target_modules: &manifest.fixed_target_modules,
+                dynamic_adapters,
+            },
+        )?));
+    }
+    Ok(stable_bytes_digest(&serde_json::to_vec(
+        &CheckpointGenerationIdentity {
+            checkpoint_generation: &manifest.checkpoint_generation,
+            step: manifest.step,
+            fixed_optimizer_step: manifest.effective_fixed_optimizer_step(),
+            model_path: &manifest.model_path,
+            lora_rank: manifest.lora_rank,
+            lora_alpha_bits: manifest.lora_alpha.to_bits(),
+            files: &manifest.files,
+            dynamic_adapters: &manifest.dynamic_adapters,
+            fixed_shard_layouts: &manifest.fixed_shard_layouts,
+            fixed_slot_identities: &manifest.fixed_slot_identities,
+        },
+    )?))
+}
+
+fn checkpoint_rank_receipt(
+    manifest: &CheckpointManifest,
+    manifest_contents: &[u8],
+) -> Result<CheckpointRankReceipt> {
+    let parallel = manifest
+        .parallel
+        .clone()
+        .context("distributed checkpoint receipt requires topology metadata")?;
+    let checkpoint_generation = manifest
+        .checkpoint_generation
+        .clone()
+        .filter(|generation| !generation.is_empty())
+        .context("distributed checkpoint receipt requires a non-empty generation")?;
+    let shard_identities = manifest
+        .tensor_shards
+        .iter()
+        .map(tensor_shard_identity)
+        .collect::<std::collections::BTreeSet<_>>();
+    let shard_identity_digest = stable_bytes_digest(&serde_json::to_vec(&shard_identities)?);
+    let all_files_declared = manifest.file_digests.len() == manifest.files.len()
+        && manifest
+            .files
+            .iter()
+            .all(|file| manifest.file_digests.contains_key(file))
+        && manifest
+            .tensor_shards
+            .iter()
+            .all(|shard| manifest.files.iter().any(|file| file == &shard.file));
+    let data_replica_metadata_complete = parallel.data_parallel_size <= 1
+        || manifest
+            .tensor_shards
+            .iter()
+            .all(|shard| shard.replicated_axes.contains(&ParallelAxis::Data));
+    let stage_state_digest = if manifest.format == STAGE_UNION_CHECKPOINT_FORMAT {
+        stable_bytes_digest(&serde_json::to_vec(&StageUnionLocalStateIdentity {
+            fixed_shard_layouts: &manifest.fixed_shard_layouts,
+            fixed_slot_identities: &manifest.fixed_slot_identities,
+            dynamic_adapters: manifest
+                .dynamic_adapters
+                .iter()
+                .map(|adapter| StageUnionLocalAdapterIdentity {
+                    id: adapter.id,
+                    shard_layouts: &adapter.shard_layouts,
+                    slot_identities: &adapter.slot_identities,
+                    parameter_count: adapter.parameter_count,
+                    optimizer_count: adapter.optimizer_count,
+                })
+                .collect(),
+        })?)
+    } else {
+        String::new()
+    };
+    Ok(CheckpointRankReceipt {
+        format: RANK_RECEIPT_FORMAT.to_string(),
+        checkpoint_format: manifest.format.clone(),
+        checkpoint_generation,
+        global_rank: parallel.global_rank,
+        parallel,
+        manifest_digest: stable_bytes_digest(manifest_contents),
+        generation_digest: checkpoint_generation_digest(manifest)?,
+        shard_identity_digest,
+        shard_count: manifest.tensor_shards.len(),
+        shard_identities_unique: shard_identities.len() == manifest.tensor_shards.len(),
+        files: manifest.files.clone(),
+        file_digests: manifest.file_digests.clone(),
+        all_files_declared,
+        data_replica_metadata_complete,
+        pipeline_stage: manifest.pipeline_stage.clone(),
+        stage_state_digest,
+    })
+}
+
+pub fn export_distributed_adapter_checkpoint(
+    final_path: &Path,
+    generation: &str,
+    parallel: &ParallelCheckpointManifest,
+    adapter_id: Option<i64>,
+    save_rank: impl FnOnce(&Path) -> Result<()>,
+) -> Result<usize> {
+    if parallel.world_size <= 1 {
+        bail!("distributed adapter export requires more than one rank");
+    }
+    if final_path.extension().is_some() {
+        bail!(
+            "distributed adapter export path must be a directory without a file extension: {}",
+            final_path.display()
+        );
+    }
+    if generation.is_empty()
+        || !generation
+            .bytes()
+            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
+    {
+        bail!("distributed adapter generation contains invalid path characters");
+    }
+    if final_path.exists() {
+        bail!(
+            "adapter export destination already exists: {}",
+            final_path.display()
+        );
+    }
+    let attempt = generation.to_string();
+    let staging = coordinate_export_attempt(final_path, generation, parallel)?;
+    let result = (|| -> Result<usize> {
+        save_rank(&staging)?;
+        wait_for_rank_manifests(&staging, parallel.world_size)?;
+        if parallel.global_rank == 0 {
+            if final_path.exists() {
+                bail!(
+                    "adapter export destination already exists: {}",
+                    final_path.display()
+                );
+            }
+            let merged = merge_distributed_adapter_checkpoint(&staging, adapter_id)?;
+            let artifact = merged_adapter_artifact(merged)?;
+            let count = artifact.tensors.len();
+            let partial_path = sibling_with_suffix(final_path, &format!(".partial-{generation}"))?;
+            if partial_path.exists() {
+                bail!(
+                    "adapter export partial destination already exists: {}",
+                    partial_path.display()
+                );
+            }
+            artifact.save(&partial_path)?;
+            write_atomic(
+                &partial_path.join(".rustrain-export-completed.json"),
+                &serde_json::to_vec_pretty(&AdapterExportCompletion {
+                    attempt: attempt.clone(),
+                    tensor_count: count,
+                })?,
+            )?;
+            std::fs::rename(&partial_path, final_path).with_context(|| {
+                format!(
+                    "publish adapter {} to {}",
+                    partial_path.display(),
+                    final_path.display()
+                )
+            })?;
+            Ok(count)
+        } else {
+            wait_for_export_completion(final_path, &attempt, &staging)
+        }
+    })();
+    if let Err(error) = &result {
+        let errors = staging.join("errors");
+        let _ = std::fs::create_dir_all(&errors);
+        let _ = write_atomic(
+            &errors.join(format!("rank-{:05}.txt", parallel.global_rank)),
+            error.to_string().as_bytes(),
+        );
+    }
+    result
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+struct AdapterExportCompletion {
+    attempt: String,
+    tensor_count: usize,
+}
+
+fn coordinate_export_attempt(
+    final_path: &Path,
+    generation: &str,
+    parallel: &ParallelCheckpointManifest,
+) -> Result<PathBuf> {
+    let staging = sibling_with_suffix(final_path, &format!(".rustrain-shards-{generation}"))?;
+    if parallel.global_rank == 0 {
+        std::fs::create_dir(&staging).with_context(|| {
+            format!(
+                "create unique distributed export staging {}; attempt IDs cannot be reused",
+                staging.display()
+            )
+        })?;
+        let result = write_atomic(&staging.join("READY"), generation.as_bytes());
+        if let Err(error) = &result {
+            let errors = staging.join("errors");
+            let _ = std::fs::create_dir_all(&errors);
+            let _ = write_atomic(&errors.join("rank-00000.txt"), error.to_string().as_bytes());
+        }
+        result?;
+    } else {
+        wait_for_distributed_export(&staging, || {
+            read_marker(&staging.join("READY")).as_deref() == Some(generation)
+        })?;
+    }
+    Ok(staging)
+}
+
+fn read_marker(path: &Path) -> Option<String> {
+    std::fs::read_to_string(path)
+        .ok()
+        .map(|value| value.trim().to_string())
+        .filter(|value| !value.is_empty())
+}
+
+pub fn merged_adapter_artifact(merged: MergedAdapterCheckpoint) -> Result<Qwen36AdapterArtifact> {
+    let runtime_config = crate::config::read_qwen36_runtime_config(Path::new(&merged.model_path))?;
+    let target_modules = merged
+        .target_modules
+        .iter()
+        .map(|module| Qwen36LoraTargetModule::parse(module))
+        .collect::<Result<Vec<_>>>()?;
+    let lora_config = Qwen36LoraConfig {
+        rank: merged.rank,
+        alpha: merged.alpha,
+        target_layers: merged.target_layers.clone(),
+        target_modules,
+    };
+    let mut by_index = merged
+        .slots
+        .into_iter()
+        .map(|slot| (slot.identity.index, slot))
+        .collect::<BTreeMap<_, _>>();
+    let mut exported = Vec::new();
+    for slot in native_lora_slots(&runtime_config, &lora_config) {
+        if !slot.active {
+            let placeholder = Tensor::zeros([], (tch::Kind::Float, tch::Device::Cpu));
+            exported.push((placeholder.shallow_clone(), placeholder));
+            continue;
+        }
+        let merged_slot = by_index
+            .remove(&slot.index)
+            .with_context(|| format!("merged adapter is missing native slot {}", slot.index))?;
+        if merged_slot.identity.layer != slot.layer
+            || merged_slot.identity.module != slot.module.cpp_name()
+        {
+            bail!(
+                "merged adapter slot {} identity is inconsistent",
+                slot.index
+            );
+        }
+        exported.push((merged_slot.lora_a, merged_slot.lora_b));
+    }
+    if !by_index.is_empty() {
+        bail!("merged adapter contains slots not present in the runtime model");
+    }
+    Qwen36AdapterArtifact::from_native_exports(
+        &merged.model_path,
+        "qwen3_hybrid_lora_sft",
+        Some(Path::new(&merged.model_path)),
+        &runtime_config,
+        &lora_config,
+        exported,
+    )
+}
+
+fn sibling_with_suffix(path: &Path, suffix: &str) -> Result<PathBuf> {
+    let file_name = path
+        .file_name()
+        .and_then(|name| name.to_str())
+        .context("adapter export path has no UTF-8 file name")?;
+    Ok(path.with_file_name(format!("{file_name}{suffix}")))
+}
+
+fn wait_for_rank_manifests(root: &Path, world_size: usize) -> Result<()> {
+    wait_for_distributed_export(root, || rank_manifests_ready(root, world_size))
+}
+
+fn rank_manifests_ready(root: &Path, world_size: usize) -> bool {
+    (0..world_size).all(|rank| {
+        let rank_dir = root.join(format!("rank-{rank:05}"));
+        rank_dir.join("manifest.json").is_file() && rank_dir.join(RANK_RECEIPT_FILE).is_file()
+    })
+}
+
+fn wait_for_export_completion(final_path: &Path, attempt: &str, staging: &Path) -> Result<usize> {
+    let completion_path = final_path.join(".rustrain-export-completed.json");
+    wait_for_distributed_export(staging, || {
+        std::fs::read(&completion_path)
+            .ok()
+            .and_then(|contents| serde_json::from_slice::<AdapterExportCompletion>(&contents).ok())
+            .is_some_and(|completion| completion.attempt == attempt)
+    })?;
+    let completion: AdapterExportCompletion = serde_json::from_slice(
+        &std::fs::read(&completion_path)
+            .with_context(|| format!("read {}", completion_path.display()))?,
+    )
+    .with_context(|| format!("parse {}", completion_path.display()))?;
+    if completion.attempt != attempt {
+        bail!("distributed adapter completion belongs to another attempt");
+    }
+    Ok(completion.tensor_count)
+}
+
+fn wait_for_distributed_export(root: &Path, ready: impl Fn() -> bool) -> Result<()> {
+    let deadline = std::time::Instant::now() + std::time::Duration::from_secs(120);
+    loop {
+        let errors = root.join("errors");
+        if errors.is_dir() {
+            let mut entries = std::fs::read_dir(&errors)?.collect::<std::io::Result<Vec<_>>>()?;
+            entries.sort_by_key(|entry| entry.file_name());
+            if let Some(error) = entries.first() {
+                let message = std::fs::read_to_string(error.path()).unwrap_or_else(|read_error| {
+                    format!("failed to read distributed export error: {read_error}")
+                });
+                bail!("distributed adapter export failed: {message}");
+            }
+        }
+        if ready() {
+            return Ok(());
+        }
+        if std::time::Instant::now() >= deadline {
+            bail!(
+                "timed out waiting for distributed adapter export at {}",
+                root.display()
+            );
+        }
+        std::thread::sleep(std::time::Duration::from_millis(20));
+    }
+}
+
+/// Load checkpoint from a directory.
+pub fn load_checkpoint(dir: &Path) -> Result<CheckpointData> {
+    load_checkpoint_at(dir, None)
+}
+
+pub fn load_checkpoint_for_topology(
+    dir: &Path,
+    parallel: &ParallelCheckpointManifest,
+) -> Result<CheckpointData> {
+    if !parallel.is_distributed() {
+        return load_checkpoint(dir);
+    }
+    let rank_dir = rank_checkpoint_dir(dir, parallel.global_rank);
+    preflight_distributed_checkpoint_set(dir, &rank_dir, parallel)?;
+    load_checkpoint_at(&rank_dir, Some(parallel))
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+struct DataReplicaKey {
+    file: String,
+    tensor_name: String,
+    state: String,
+    adapter_id: Option<i64>,
+    tensor_rank: usize,
+    pipeline_rank: usize,
+    expert_rank: usize,
+    context_rank: usize,
+}
+
+struct DataReplica {
+    global_rank: usize,
+    shard: TensorShardManifest,
+    tensor: Tensor,
+}
+
+struct DataReplicaGroup {
+    data_ranks: std::collections::BTreeSet<usize>,
+    reference: DataReplica,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+struct DataReplicaFileKey {
+    file: String,
+    tensor_rank: usize,
+    pipeline_rank: usize,
+    expert_rank: usize,
+    context_rank: usize,
+}
+
+struct DataReplicaFileGroup {
+    data_ranks: std::collections::BTreeSet<usize>,
+    global_rank: usize,
+    digest: String,
+}
+
+fn read_checkpoint_manifest(path: &Path) -> Result<CheckpointManifest> {
+    serde_json::from_str(
+        &std::fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?,
+    )
+    .with_context(|| format!("parse {}", path.display()))
+}
+
+fn read_checkpoint_rank_receipt(path: &Path) -> Result<CheckpointRankReceipt> {
+    serde_json::from_str(
+        &std::fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?,
+    )
+    .with_context(|| format!("parse {}", path.display()))
+}
+
+fn preflight_distributed_receipts(
+    root: &Path,
+    requested_rank_dir: &Path,
+    requested_manifest: &CheckpointManifest,
+    expected: &ParallelCheckpointManifest,
+    topology: &ParallelTopology,
+) -> Result<()> {
+    let requested_manifest_path = requested_rank_dir.join("manifest.json");
+    let requested_manifest_contents = std::fs::read(&requested_manifest_path)
+        .with_context(|| format!("read {}", requested_manifest_path.display()))?;
+    let expected_local_receipt =
+        checkpoint_rank_receipt(requested_manifest, &requested_manifest_contents)?;
+    let stage_union = requested_manifest.format == STAGE_UNION_CHECKPOINT_FORMAT;
+    let mut baseline: Option<CheckpointRankReceipt> = None;
+    let mut stage_baselines: BTreeMap<usize, CheckpointRankReceipt> = BTreeMap::new();
+    let mut replica_files: BTreeMap<DataReplicaFileKey, DataReplicaFileGroup> = BTreeMap::new();
+
+    for global_rank in 0..expected.world_size {
+        let rank_dir = rank_checkpoint_dir(root, global_rank);
+        let manifest_path = rank_dir.join("manifest.json");
+        if !manifest_path.is_file() {
+            bail!(
+                "distributed checkpoint is missing rank {global_rank} manifest {}",
+                manifest_path.display()
+            );
+        }
+        let receipt_path = rank_dir.join(RANK_RECEIPT_FILE);
+        let receipt = read_checkpoint_rank_receipt(&receipt_path)?;
+        let expected_parallel =
+            ParallelCheckpointManifest::from_topology(expected.world_size, global_rank, topology)?;
+        if receipt.format != RANK_RECEIPT_FORMAT
+            || receipt.checkpoint_format != requested_manifest.format
+            || receipt.global_rank != global_rank
+            || receipt.parallel != expected_parallel
+        {
+            bail!(
+                "distributed checkpoint topology mismatch: rank {global_rank} receipt format or topology is inconsistent"
+            );
+        }
+        if !receipt.shard_identities_unique {
+            bail!("rank {global_rank} checkpoint contains duplicate tensor shard identities");
+        }
+        if !receipt.all_files_declared {
+            bail!("rank {global_rank} checkpoint receipt has incomplete file declarations");
+        }
+        if !receipt.data_replica_metadata_complete {
+            bail!("rank {global_rank} checkpoint is missing data-parallel replica metadata");
+        }
+        if receipt.file_digests.len() != receipt.files.len()
+            || receipt
+                .files
+                .iter()
+                .any(|file| !receipt.file_digests.contains_key(file))
+        {
+            bail!("rank {global_rank} checkpoint receipt is missing file content digests");
+        }
+        if global_rank == expected.global_rank && receipt != expected_local_receipt {
+            bail!("local checkpoint manifest does not match its compact rank receipt");
+        }
+
+        if let Some(baseline) = &baseline {
+            if receipt.checkpoint_generation != baseline.checkpoint_generation
+                || receipt.generation_digest != baseline.generation_digest
+                || receipt.files != baseline.files
+            {
+                bail!("rank {global_rank} receipt belongs to a different checkpoint generation");
+            }
+            if !stage_union
+                && (receipt.shard_identity_digest != baseline.shard_identity_digest
+                    || receipt.shard_count != baseline.shard_count)
+            {
+                bail!("rank {global_rank} receipt belongs to a different checkpoint generation");
+            }
+        } else {
+            baseline = Some(receipt.clone());
+        }
+
+        if stage_union {
+            let pipeline_rank = receipt.parallel.coordinates.pipeline;
+            if let Some(stage_baseline) = stage_baselines.get(&pipeline_rank) {
+                if receipt.pipeline_stage != stage_baseline.pipeline_stage
+                    || receipt.stage_state_digest != stage_baseline.stage_state_digest
+                    || receipt.shard_identity_digest != stage_baseline.shard_identity_digest
+                    || receipt.shard_count != stage_baseline.shard_count
+                {
+                    bail!(
+                        "rank {global_rank} stage-local checkpoint inventory differs within pipeline stage {pipeline_rank}"
+                    );
+                }
+            } else {
+                stage_baselines.insert(pipeline_rank, receipt.clone());
+            }
+        }
+
+        if expected.data_parallel_size <= 1 {
+            continue;
+        }
+        for file in &receipt.files {
+            let digest = receipt
+                .file_digests
+                .get(file)
+                .context("validated checkpoint receipt digest disappeared")?;
+            let coordinates = &receipt.parallel.coordinates;
+            let key = DataReplicaFileKey {
+                file: file.clone(),
+                tensor_rank: coordinates.tensor,
+                pipeline_rank: coordinates.pipeline,
+                expert_rank: coordinates.expert,
+                context_rank: coordinates.context,
+            };
+            match replica_files.entry(key) {
+                std::collections::btree_map::Entry::Vacant(entry) => {
+                    entry.insert(DataReplicaFileGroup {
+                        data_ranks: [coordinates.data].into_iter().collect(),
+                        global_rank,
+                        digest: digest.clone(),
+                    });
+                }
+                std::collections::btree_map::Entry::Occupied(mut entry) => {
+                    let key = entry.key().clone();
+                    let group = entry.get_mut();
+                    if !group.data_ranks.insert(coordinates.data) {
+                        bail!("duplicate data-parallel checkpoint replica at rank {global_rank}");
+                    }
+                    if group.digest != *digest {
+                        bail!(
+                            "data-parallel replica content digest differs for {} between ranks {} and {}",
+                            key.file,
+                            group.global_rank,
+                            global_rank
+                        );
+                    }
+                }
+            }
+        }
+    }
+
+    for (key, group) in replica_files {
+        if group.data_ranks.len() != expected.data_parallel_size {
+            bail!(
+                "data-parallel checkpoint replica set for {} has {} ranks, expected {}",
+                key.file,
+                group.data_ranks.len(),
+                expected.data_parallel_size
+            );
+        }
+    }
+    if stage_union {
+        validate_stage_union_receipt_coverage(&stage_baselines, expected)?;
+    }
+    Ok(())
+}
+
+fn validate_stage_union_receipt_coverage(
+    stages: &BTreeMap<usize, CheckpointRankReceipt>,
+    expected: &ParallelCheckpointManifest,
+) -> Result<()> {
+    if stages.len() != expected.pipeline_model_parallel_size {
+        bail!(
+            "stage-union checkpoint has {} pipeline stages, expected {}",
+            stages.len(),
+            expected.pipeline_model_parallel_size
+        );
+    }
+    let mut next_layer = 0usize;
+    let mut global_num_layers = None;
+    for pipeline_rank in 0..expected.pipeline_model_parallel_size {
+        let stage = stages
+            .get(&pipeline_rank)
+            .and_then(|receipt| receipt.pipeline_stage.as_ref())
+            .with_context(|| {
+                format!("stage-union checkpoint is missing pipeline stage {pipeline_rank}")
+            })?;
+        if stage.pipeline_rank != pipeline_rank
+            || stage.pipeline_size != expected.pipeline_model_parallel_size
+            || stage.layer_start != next_layer
+        {
+            bail!(
+                "stage-union checkpoint has invalid or non-contiguous metadata for pipeline stage {pipeline_rank}"
+            );
+        }
+        if let Some(expected_layers) = global_num_layers {
+            if stage.global_num_layers != expected_layers {
+                bail!("stage-union checkpoint stages disagree on global layer count");
+            }
+        } else {
+            global_num_layers = Some(stage.global_num_layers);
+        }
+        next_layer = stage.layer_end;
+    }
+    if Some(next_layer) != global_num_layers {
+        bail!("stage-union checkpoint layer ranges do not cover the global model");
+    }
+    Ok(())
+}
+
+/// Validate a complete v5 rank set before restoring rank-local state.
+///
+/// New v5 manifests bind tensor files to content digests. Receipt-aware
+/// checkpoints read compact metadata for every rank and validate only local
+/// full state. V5 checkpoints that predate receipts retain the full-manifest
+/// compatibility preflight; digestless v5 also reads tensors for DP parity.
+fn preflight_distributed_checkpoint_set(
+    root: &Path,
+    requested_rank_dir: &Path,
+    expected: &ParallelCheckpointManifest,
+) -> Result<()> {
+    let requested_manifest = read_checkpoint_manifest(&requested_rank_dir.join("manifest.json"))?;
+    if !is_current_distributed_checkpoint(&requested_manifest.format) {
+        // V3/V4 rank shards predate complete-set and DP replica metadata.
+        return Ok(());
+    }
+
+    let rank_order = expected
+        .rank_order
+        .iter()
+        .map(|axis| axis.name())
+        .collect::<Vec<_>>()
+        .join("-");
+    let topology = ParallelTopology::with_order(
+        expected.tensor_model_parallel_size,
+        expected.pipeline_model_parallel_size,
+        expected.data_parallel_size,
+        expected.expert_model_parallel_size,
+        expected.context_parallel_size,
+        &rank_order,
+    )?;
+    topology.validate_world_size(expected.world_size)?;
+    let expected_requested = ParallelCheckpointManifest::from_topology(
+        expected.world_size,
+        expected.global_rank,
+        &topology,
+    )?;
+    if &expected_requested != expected {
+        bail!(
+            "distributed checkpoint topology mismatch: current topology metadata is internally inconsistent"
+        );
+    }
+
+    match requested_manifest.rank_receipt_version {
+        Some(RANK_RECEIPT_VERSION) => {
+            return preflight_distributed_receipts(
+                root,
+                requested_rank_dir,
+                &requested_manifest,
+                expected,
+                &topology,
+            );
+        }
+        Some(version) => {
+            bail!("unsupported distributed checkpoint rank receipt version {version}");
+        }
+        None if requested_manifest.format == STAGE_UNION_CHECKPOINT_FORMAT => {
+            bail!("stage-union checkpoints require compact rank receipts");
+        }
+        None => {}
+    }
+
+    let rank_zero_path = rank_checkpoint_dir(root, 0).join("manifest.json");
+    let rank_zero = read_checkpoint_manifest(&rank_zero_path)?;
+    if rank_zero.format != TP_CHECKPOINT_FORMAT {
+        bail!("distributed checkpoint rank set mixes v5 and legacy checkpoint formats");
+    }
+    let rank_zero_parallel = rank_zero
+        .parallel
+        .as_ref()
+        .context("rank 0 v5 checkpoint is missing topology metadata")?;
+    let expected_rank_zero =
+        ParallelCheckpointManifest::from_topology(expected.world_size, 0, &topology)?;
+    if rank_zero_parallel != &expected_rank_zero {
+        bail!(
+            "distributed checkpoint topology mismatch: rank 0 does not match the current topology"
+        );
+    }
+
+    let baseline_shards = rank_zero
+        .tensor_shards
+        .iter()
+        .map(tensor_shard_identity)
+        .collect::<std::collections::BTreeSet<_>>();
+    if baseline_shards.len() != rank_zero.tensor_shards.len() {
+        bail!("rank 0 checkpoint contains duplicate tensor shard identities");
+    }
+
+    let digest_preflight = !rank_zero.file_digests.is_empty();
+    let mut replicas: BTreeMap<DataReplicaKey, DataReplicaGroup> = BTreeMap::new();
+    let mut replica_files: BTreeMap<DataReplicaFileKey, DataReplicaFileGroup> = BTreeMap::new();
+    for global_rank in 0..expected.world_size {
+        let rank_dir = rank_checkpoint_dir(root, global_rank);
+        let manifest_path = rank_dir.join("manifest.json");
+        let manifest = if global_rank == expected.global_rank {
+            requested_manifest.clone()
+        } else if global_rank == 0 {
+            rank_zero.clone()
+        } else {
+            read_checkpoint_manifest(&manifest_path)?
+        };
+        let saved_parallel = manifest
+            .parallel
+            .as_ref()
+            .with_context(|| format!("rank {global_rank} v5 checkpoint is missing topology"))?;
+        let expected_parallel =
+            ParallelCheckpointManifest::from_topology(expected.world_size, global_rank, &topology)?;
+        if manifest.format != TP_CHECKPOINT_FORMAT || saved_parallel != &expected_parallel {
+            bail!(
+                "distributed checkpoint topology mismatch: rank {global_rank} topology or format is inconsistent with the complete v5 rank set"
+            );
+        }
+        validate_checkpoint_generation(&rank_zero, &manifest, global_rank)?;
+
+        let shard_identities = manifest
+            .tensor_shards
+            .iter()
+            .map(tensor_shard_identity)
+            .collect::<std::collections::BTreeSet<_>>();
+        if shard_identities.len() != manifest.tensor_shards.len()
+            || shard_identities != baseline_shards
+        {
+            bail!("rank {global_rank} checkpoint tensor shard set differs from rank 0");
+        }
+
+        if digest_preflight {
+            if manifest.file_digests.len() != manifest.files.len()
+                || manifest
+                    .files
+                    .iter()
+                    .any(|file| !manifest.file_digests.contains_key(file))
+            {
+                bail!("rank {global_rank} checkpoint is missing tensor file content digests");
+            }
+            for shard in &manifest.tensor_shards {
+                if !manifest.files.iter().any(|saved| saved == &shard.file) {
+                    bail!(
+                        "rank {global_rank} tensor shard references undeclared file {}",
+                        shard.file
+                    );
+                }
+                if expected.data_parallel_size > 1
+                    && !shard.replicated_axes.contains(&ParallelAxis::Data)
+                {
+                    bail!(
+                        "rank {global_rank} tensor {} is missing data-parallel replica metadata",
+                        shard.tensor_name
+                    );
+                }
+            }
+            for file in &manifest.files {
+                let digest = manifest
+                    .file_digests
+                    .get(file)
+                    .context("validated checkpoint file digest disappeared")?;
+                if expected.data_parallel_size <= 1 {
+                    continue;
+                }
+                let key = DataReplicaFileKey {
+                    file: file.clone(),
+                    tensor_rank: saved_parallel.coordinates.tensor,
+                    pipeline_rank: saved_parallel.coordinates.pipeline,
+                    expert_rank: saved_parallel.coordinates.expert,
+                    context_rank: saved_parallel.coordinates.context,
+                };
+                match replica_files.entry(key) {
+                    std::collections::btree_map::Entry::Vacant(entry) => {
+                        entry.insert(DataReplicaFileGroup {
+                            data_ranks: [saved_parallel.coordinates.data].into_iter().collect(),
+                            global_rank,
+                            digest: digest.clone(),
+                        });
+                    }
+                    std::collections::btree_map::Entry::Occupied(mut entry) => {
+                        let key = entry.key().clone();
+                        let group = entry.get_mut();
+                        if !group.data_ranks.insert(saved_parallel.coordinates.data) {
+                            bail!(
+                                "duplicate data-parallel checkpoint replica at rank {global_rank}"
+                            );
+                        }
+                        if group.digest != *digest {
+                            bail!(
+                                "data-parallel replica content digest differs for {} between ranks {} and {}",
+                                key.file,
+                                group.global_rank,
+                                global_rank
+                            );
+                        }
+                    }
+                }
+            }
+            continue;
+        }
+
+        let mut tensors_by_file = BTreeMap::new();
+        for file in manifest
+            .tensor_shards
+            .iter()
+            .map(|shard| shard.file.as_str())
+            .collect::<std::collections::BTreeSet<_>>()
+        {
+            if !manifest.files.iter().any(|saved| saved == file) {
+                bail!("rank {global_rank} tensor shard references undeclared file {file}");
+            }
+            tensors_by_file.insert(file.to_string(), read_named_tensors(&rank_dir.join(file))?);
+        }
+
+        for shard in &manifest.tensor_shards {
+            let tensor = tensors_by_file
+                .get(&shard.file)
+                .and_then(|tensors| tensors.get(&shard.tensor_name))
+                .with_context(|| {
+                    format!(
+                        "rank {global_rank} checkpoint is missing tensor {}:{}",
+                        shard.file, shard.tensor_name
+                    )
+                })?;
+            if tensor.size() != shard.local_shape {
+                bail!(
+                    "rank {global_rank} tensor {}:{} shape {:?} does not match manifest {:?}",
+                    shard.file,
+                    shard.tensor_name,
+                    tensor.size(),
+                    shard.local_shape
+                );
+            }
+            if expected.data_parallel_size > 1
+                && !shard.replicated_axes.contains(&ParallelAxis::Data)
+            {
+                bail!(
+                    "rank {global_rank} tensor {} is missing data-parallel replica metadata",
+                    shard.tensor_name
+                );
+            }
+            if !shard.replicated_axes.contains(&ParallelAxis::Data) {
+                continue;
+            }
+            let key = DataReplicaKey {
+                file: shard.file.clone(),
+                tensor_name: shard.tensor_name.clone(),
+                state: shard.state.clone(),
+                adapter_id: shard.adapter_id,
+                tensor_rank: saved_parallel.coordinates.tensor,
+                pipeline_rank: saved_parallel.coordinates.pipeline,
+                expert_rank: saved_parallel.coordinates.expert,
+                context_rank: saved_parallel.coordinates.context,
+            };
+            let mut replica_shard = shard.clone();
+            replica_shard.replica_identity.clear();
+            let replica = DataReplica {
+                global_rank,
+                shard: replica_shard,
+                tensor: tensor.shallow_clone(),
+            };
+            match replicas.entry(key) {
+                std::collections::btree_map::Entry::Vacant(entry) => {
+                    entry.insert(DataReplicaGroup {
+                        data_ranks: [saved_parallel.coordinates.data].into_iter().collect(),
+                        reference: replica,
+                    });
+                }
+                std::collections::btree_map::Entry::Occupied(mut entry) => {
+                    let key = entry.key().clone();
+                    let group = entry.get_mut();
+                    if !group.data_ranks.insert(saved_parallel.coordinates.data) {
+                        bail!("duplicate data-parallel checkpoint replica at rank {global_rank}");
+                    }
+                    if replica.shard != group.reference.shard {
+                        bail!(
+                            "data-parallel replica metadata differs for {}:{} between ranks {} and {}",
+                            key.file,
+                            key.tensor_name,
+                            group.reference.global_rank,
+                            replica.global_rank
+                        );
+                    }
+                    if !group
+                        .reference
+                        .tensor
+                        .allclose(&replica.tensor, 0.0, 0.0, false)
+                    {
+                        bail!(
+                            "data-parallel replica tensor differs for {}:{} ({}) between ranks {} and {}",
+                            key.file,
+                            key.tensor_name,
+                            key.state,
+                            group.reference.global_rank,
+                            replica.global_rank
+                        );
+                    }
+                }
+            }
+        }
+    }
+
+    for (key, group) in replicas {
+        if group.data_ranks.len() != expected.data_parallel_size {
+            bail!(
+                "data-parallel checkpoint replica set for {}:{} has {} ranks, expected {}",
+                key.file,
+                key.tensor_name,
+                group.data_ranks.len(),
+                expected.data_parallel_size
+            );
+        }
+    }
+    for (key, group) in replica_files {
+        if group.data_ranks.len() != expected.data_parallel_size {
+            bail!(
+                "data-parallel checkpoint replica set for {} has {} ranks, expected {}",
+                key.file,
+                group.data_ranks.len(),
+                expected.data_parallel_size
+            );
+        }
+    }
+    Ok(())
+}
+
+fn tensor_shard_identity(shard: &TensorShardManifest) -> (String, String, String, Option<i64>) {
+    (
+        shard.file.clone(),
+        shard.tensor_name.clone(),
+        shard.state.clone(),
+        shard.adapter_id,
+    )
+}
+
+fn validate_checkpoint_generation(
+    baseline: &CheckpointManifest,
+    manifest: &CheckpointManifest,
+    global_rank: usize,
+) -> Result<()> {
+    if baseline.format == STAGE_UNION_CHECKPOINT_FORMAT
+        && manifest.format == STAGE_UNION_CHECKPOINT_FORMAT
+    {
+        if checkpoint_generation_digest(baseline)? != checkpoint_generation_digest(manifest)? {
+            bail!("rank {global_rank} checkpoint belongs to a different checkpoint generation");
+        }
+        return Ok(());
+    }
+    // Distributed ranks may report different local losses for the same step.
+    if manifest.checkpoint_generation != baseline.checkpoint_generation
+        || manifest.step != baseline.step
+        || manifest.effective_fixed_optimizer_step() != baseline.effective_fixed_optimizer_step()
+        || manifest.model_path != baseline.model_path
+        || manifest.lora_rank != baseline.lora_rank
+        || manifest.lora_alpha.to_bits() != baseline.lora_alpha.to_bits()
+        || manifest.files != baseline.files
+        || manifest.dynamic_adapters != baseline.dynamic_adapters
+        || manifest.fixed_shard_layouts != baseline.fixed_shard_layouts
+        || manifest.fixed_slot_identities != baseline.fixed_slot_identities
+    {
+        bail!("rank {global_rank} checkpoint belongs to a different checkpoint generation");
+    }
+    Ok(())
+}
+
+fn load_checkpoint_at(
+    dir: &Path,
+    expected_parallel: Option<&ParallelCheckpointManifest>,
+) -> Result<CheckpointData> {
+    let manifest_path = dir.join("manifest.json");
+    let manifest: CheckpointManifest = serde_json::from_str(
+        &std::fs::read_to_string(&manifest_path)
+            .with_context(|| format!("read {}", manifest_path.display()))?,
+    )
+    .with_context(|| "parse manifest.json")?;
+    validate_manifest_file_digests(dir, &manifest)?;
+    match expected_parallel {
+        Some(expected) => {
+            if !is_current_distributed_checkpoint(&manifest.format)
+                && manifest.format != PROJECTION_AWARE_TP_CHECKPOINT_FORMAT
+                && manifest.format != LEGACY_TP_CHECKPOINT_FORMAT
+            {
+                bail!(
+                    "distributed resume requires {STAGE_UNION_CHECKPOINT_FORMAT}, {TP_CHECKPOINT_FORMAT}, {PROJECTION_AWARE_TP_CHECKPOINT_FORMAT}, or {LEGACY_TP_CHECKPOINT_FORMAT}, found {}",
+                    manifest.format
+                );
+            }
+            if !is_current_distributed_checkpoint(&manifest.format)
+                && (expected.pipeline_model_parallel_size > 1
+                    || expected.data_parallel_size > 1
+                    || expected.expert_model_parallel_size > 1
+                    || expected.context_parallel_size > 1)
+            {
+                bail!(
+                    "legacy v3/v4 checkpoints only support tensor-parallel rank shards; save a v5 checkpoint for this distributed topology"
+                );
+            }
+            let saved = manifest
+                .parallel
+                .as_ref()
+                .context("tensor-parallel checkpoint is missing topology metadata")?;
+            let topology_matches = if is_current_distributed_checkpoint(&manifest.format) {
+                saved == expected
+            } else {
+                saved.legacy_fields_match(expected)
+            };
+            if !topology_matches {
+                bail!(
+                    "distributed checkpoint topology mismatch: saved={saved:?}, current={expected:?}"
+                );
+            }
+        }
+        None if is_current_distributed_checkpoint(&manifest.format)
+            || manifest.format == PROJECTION_AWARE_TP_CHECKPOINT_FORMAT
+            || manifest.format == LEGACY_TP_CHECKPOINT_FORMAT =>
+        {
+            bail!("distributed checkpoint must be loaded with rank topology");
+        }
+        None => {}
+    }
+
+    let adapter_path = dir.join("adapter.safetensors");
+    let adapter_named = read_named_tensors(&adapter_path)?;
+    let lora_a = collect_side(&adapter_named, "a_")?;
+    let lora_b = collect_side(&adapter_named, "b_")?;
+
+    let optimizer_path = dir.join("optimizer.safetensors");
+    let optimizer_named = read_named_tensors(&optimizer_path)?;
+    let adam_m = collect_side(&optimizer_named, "a_")?;
+    let adam_v = collect_side(&optimizer_named, "b_")?;
+    let mut dynamic_adapters = Vec::with_capacity(manifest.dynamic_adapters.len());
+    for dynamic_manifest in &manifest.dynamic_adapters {
+        let id = dynamic_manifest.id;
+        let dynamic_lora_a = collect_side(&adapter_named, &format!("dynamic_{id}_a_"))?;
+        let dynamic_lora_b = collect_side(&adapter_named, &format!("dynamic_{id}_b_"))?;
+        let dynamic_adam_m = collect_side(&optimizer_named, &format!("dynamic_{id}_a_"))?;
+        let dynamic_adam_v = collect_side(&optimizer_named, &format!("dynamic_{id}_b_"))?;
+        if dynamic_lora_a.len() != dynamic_manifest.parameter_count
+            || dynamic_lora_b.len() != dynamic_manifest.parameter_count
+            || dynamic_adam_m.len() != dynamic_manifest.optimizer_count
+            || dynamic_adam_v.len() != dynamic_manifest.optimizer_count
+        {
+            anyhow::bail!("dynamic adapter {id} checkpoint tensor count mismatch");
+        }
+        dynamic_adapters.push(DynamicAdapterCheckpoint {
+            manifest: dynamic_manifest.clone(),
+            lora_a: dynamic_lora_a,
+            lora_b: dynamic_lora_b,
+            adam_m: dynamic_adam_m,
+            adam_v: dynamic_adam_v,
+        });
+    }
+    validate_optimizer_clock_state(&manifest, &adam_m, &adam_v, &dynamic_adapters)?;
+
+    if let Some(parallel) = expected_parallel {
+        let mut expected_shards = build_tensor_shard_manifest(
+            parallel,
+            manifest.lora_rank,
+            &lora_a,
+            &lora_b,
+            &adam_m,
+            &adam_v,
+            &dynamic_adapters,
+            &manifest.fixed_shard_layouts,
+        )?;
+        if !is_current_distributed_checkpoint(&manifest.format) {
+            for shard in &mut expected_shards {
+                shard.placements.clear();
+                shard.replicated_axes.clear();
+            }
+        }
+        validate_saved_shards(&manifest.tensor_shards, &expected_shards)?;
+    }
+
+    tracing::info!(
+        step = manifest.step,
+        loss = manifest.loss,
+        "checkpoint loaded"
+    );
+
+    Ok(CheckpointData {
+        manifest,
+        lora_a,
+        lora_b,
+        adam_m,
+        adam_v,
+        dynamic_adapters,
+    })
+}
+
+fn validate_optimizer_clock_state(
+    manifest: &CheckpointManifest,
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+    dynamic_adapters: &[DynamicAdapterCheckpoint],
+) -> Result<()> {
+    if manifest.effective_fixed_optimizer_step() > 0
+        && (manifest.format != STAGE_UNION_CHECKPOINT_FORMAT
+            || !manifest.fixed_slot_identities.is_empty())
+        && (adam_m.is_empty() || adam_v.is_empty())
+    {
+        bail!(
+            "checkpoint fixed optimizer step {} has no Adam state",
+            manifest.effective_fixed_optimizer_step()
+        );
+    }
+    for adapter in dynamic_adapters {
+        if adapter.manifest.optimizer_step > 0
+            && (manifest.format != STAGE_UNION_CHECKPOINT_FORMAT
+                || adapter.manifest.parameter_count > 0)
+            && (adapter.adam_m.is_empty() || adapter.adam_v.is_empty())
+        {
+            bail!(
+                "dynamic adapter {} optimizer step {} has no Adam state",
+                adapter.manifest.id,
+                adapter.manifest.optimizer_step
+            );
+        }
+    }
+    Ok(())
+}
+
+fn validate_manifest_file_digests(dir: &Path, manifest: &CheckpointManifest) -> Result<()> {
+    if manifest.file_digests.is_empty() {
+        return Ok(());
+    }
+    if manifest.file_digests.len() != manifest.files.len()
+        || manifest
+            .files
+            .iter()
+            .any(|file| !manifest.file_digests.contains_key(file))
+    {
+        bail!("checkpoint manifest has an incomplete tensor file digest set");
+    }
+    for file in &manifest.files {
+        let expected = manifest
+            .file_digests
+            .get(file)
+            .context("validated checkpoint file digest disappeared")?;
+        let actual = stable_file_digest(&dir.join(file))?;
+        if actual != *expected {
+            bail!("checkpoint content digest differs from manifest for {file}");
+        }
+    }
+    Ok(())
+}
+
+#[derive(Clone, Copy)]
+enum LoraSide {
+    A,
+    B,
+}
+
+fn validate_tensor_counts(
+    lora_a: &[Tensor],
+    lora_b: &[Tensor],
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+    dynamic_adapters: &[DynamicAdapterCheckpoint],
+    strict_optimizer_layout: bool,
+) -> Result<()> {
+    if lora_a.len() != lora_b.len() || adam_m.len() != adam_v.len() {
+        bail!("fixed adapter checkpoint tensor count mismatch");
+    }
+    if strict_optimizer_layout
+        && !adam_m.is_empty()
+        && adam_m.len() != lora_a.len().saturating_mul(2)
+    {
+        bail!("tensor-parallel fixed optimizer state must contain A/B entries for every LoRA slot");
+    }
+    for adapter in dynamic_adapters {
+        if adapter
+            .manifest
+            .optimizer_lr
+            .is_some_and(|lr| !lr.is_finite() || lr < 0.0)
+        {
+            bail!(
+                "dynamic adapter {} optimizer learning rate must be finite and non-negative",
+                adapter.manifest.id
+            );
+        }
+        for (name, value) in [
+            ("beta1", adapter.manifest.optimizer_beta1),
+            ("beta2", adapter.manifest.optimizer_beta2),
+        ] {
+            if value.is_some_and(|value| !value.is_finite() || !(0.0..1.0).contains(&value)) {
+                bail!(
+                    "dynamic adapter {} optimizer {name} must be finite and in [0, 1)",
+                    adapter.manifest.id
+                );
+            }
+        }
+        if adapter
+            .manifest
+            .optimizer_eps
+            .is_some_and(|value| !value.is_finite() || value < 0.0)
+        {
+            bail!(
+                "dynamic adapter {} optimizer epsilon must be finite and non-negative",
+                adapter.manifest.id
+            );
+        }
+        if adapter.lora_a.len() != adapter.lora_b.len()
+            || adapter.adam_m.len() != adapter.adam_v.len()
+            || adapter.manifest.parameter_count != adapter.lora_a.len()
+            || adapter.manifest.optimizer_count != adapter.adam_m.len()
+        {
+            bail!(
+                "dynamic adapter {} checkpoint count mismatch",
+                adapter.manifest.id
+            );
+        }
+        if strict_optimizer_layout
+            && !adapter.adam_m.is_empty()
+            && adapter.adam_m.len() != adapter.lora_a.len().saturating_mul(2)
+        {
+            bail!(
+                "tensor-parallel dynamic adapter {} optimizer state must contain A/B entries for every LoRA slot",
+                adapter.manifest.id
+            );
+        }
+    }
+    Ok(())
+}
+
+fn build_tensor_shard_manifest(
+    parallel: &ParallelCheckpointManifest,
+    fixed_lora_rank: i64,
+    lora_a: &[Tensor],
+    lora_b: &[Tensor],
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+    dynamic_adapters: &[DynamicAdapterCheckpoint],
+    fixed_shard_layouts: &[LoraTpShardLayout],
+) -> Result<Vec<TensorShardManifest>> {
+    validate_tensor_counts(lora_a, lora_b, adam_m, adam_v, dynamic_adapters, true)?;
+    let mut shards = Vec::new();
+    append_adapter_shards(
+        &mut shards,
+        parallel,
+        None,
+        fixed_lora_rank,
+        lora_a,
+        lora_b,
+        adam_m,
+        adam_v,
+        fixed_shard_layouts,
+    )?;
+    for adapter in dynamic_adapters {
+        append_adapter_shards(
+            &mut shards,
+            parallel,
+            Some(adapter.manifest.id),
+            adapter.manifest.rank,
+            &adapter.lora_a,
+            &adapter.lora_b,
+            &adapter.adam_m,
+            &adapter.adam_v,
+            &adapter.manifest.shard_layouts,
+        )?;
+    }
+    Ok(shards)
+}
+
+#[allow(clippy::too_many_arguments)]
+fn append_adapter_shards(
+    shards: &mut Vec<TensorShardManifest>,
+    parallel: &ParallelCheckpointManifest,
+    adapter_id: Option<i64>,
+    global_lora_rank: i64,
+    lora_a: &[Tensor],
+    lora_b: &[Tensor],
+    adam_m: &[Tensor],
+    adam_v: &[Tensor],
+    shard_layouts: &[LoraTpShardLayout],
+) -> Result<()> {
+    if !shard_layouts.is_empty() && shard_layouts.len() != lora_a.len() {
+        bail!(
+            "LoRA shard layout count {} does not match parameter count {}",
+            shard_layouts.len(),
+            lora_a.len()
+        );
+    }
+    let layout = |index: usize| {
+        shard_layouts
+            .get(index)
+            .copied()
+            .unwrap_or(LoraTpShardLayout::LatentRank)
+    };
+    let adapter_prefix = adapter_id
+        .map(|id| format!("dynamic_{id}_"))
+        .unwrap_or_default();
+    for (index, tensor) in lora_a.iter().enumerate() {
+        shards.push(tensor_shard(
+            parallel,
+            adapter_id,
+            global_lora_rank,
+            "adapter.safetensors",
+            format!("{adapter_prefix}a_{index}"),
+            "lora_a",
+            LoraSide::A,
+            layout(index),
+            tensor,
+        )?);
+    }
+    for (index, tensor) in lora_b.iter().enumerate() {
+        shards.push(tensor_shard(
+            parallel,
+            adapter_id,
+            global_lora_rank,
+            "adapter.safetensors",
+            format!("{adapter_prefix}b_{index}"),
+            "lora_b",
+            LoraSide::B,
+            layout(index),
+            tensor,
+        )?);
+    }
+    for (index, tensor) in adam_m.iter().enumerate() {
+        shards.push(tensor_shard(
+            parallel,
+            adapter_id,
+            global_lora_rank,
+            "optimizer.safetensors",
+            format!("{adapter_prefix}a_{index}"),
+            "adam_m",
+            if index % 2 == 0 {
+                LoraSide::A
+            } else {
+                LoraSide::B
+            },
+            layout(index / 2),
+            tensor,
+        )?);
+    }
+    for (index, tensor) in adam_v.iter().enumerate() {
+        shards.push(tensor_shard(
+            parallel,
+            adapter_id,
+            global_lora_rank,
+            "optimizer.safetensors",
+            format!("{adapter_prefix}b_{index}"),
+            "adam_v",
+            if index % 2 == 0 {
+                LoraSide::A
+            } else {
+                LoraSide::B
+            },
+            layout(index / 2),
+            tensor,
+        )?);
+    }
+    Ok(())
+}
+
+#[allow(clippy::too_many_arguments)]
+fn tensor_shard(
+    parallel: &ParallelCheckpointManifest,
+    adapter_id: Option<i64>,
+    global_lora_rank: i64,
+    file: &str,
+    tensor_name: String,
+    state: &str,
+    side: LoraSide,
+    layout: LoraTpShardLayout,
+    tensor: &Tensor,
+) -> Result<TensorShardManifest> {
+    if parallel.context_parallel_size > 1 {
+        bail!(
+            "checkpoint does not yet support context-sharded LoRA state (cp={})",
+            parallel.context_parallel_size
+        );
+    }
+    let tp_size = i64::try_from(parallel.tensor_model_parallel_size)
+        .context("TP size exceeds checkpoint tensor shape range")?;
+    let tp_rank = i64::try_from(parallel.tensor_model_parallel_rank)
+        .context("TP rank exceeds checkpoint tensor shape range")?;
+    if global_lora_rank <= 0 {
+        bail!("global LoRA rank must be positive");
+    }
+    if layout == LoraTpShardLayout::LatentRank && global_lora_rank % tp_size != 0 {
+        bail!(
+            "latent-rank sharded LoRA rank {global_lora_rank} must be divisible by TP size {tp_size}"
+        );
+    }
+    let local_shape = tensor.size();
+    if local_shape.len() < 2 {
+        bail!("checkpoint tensor {file}:{tensor_name} must have at least two dimensions");
+    }
+    let rank_axis = match side {
+        LoraSide::A => local_shape.len() - 2,
+        LoraSide::B => local_shape.len() - 1,
+    };
+    let local_lora_rank = global_lora_rank / tp_size;
+    let mut global_shape = local_shape.clone();
+    let mut global_offset = vec![0; local_shape.len()];
+    let (partition_axis, replicated) = match (layout, side) {
+        (LoraTpShardLayout::LatentRank, _) => {
+            if local_shape[rank_axis] != local_lora_rank {
+                bail!(
+                    "checkpoint tensor {file}:{tensor_name} has local rank {} on axis {rank_axis}, expected {local_lora_rank}",
+                    local_shape[rank_axis]
+                );
+            }
+            global_shape[rank_axis] = global_lora_rank;
+            global_offset[rank_axis] = tp_rank * local_lora_rank;
+            (rank_axis, false)
+        }
+        (LoraTpShardLayout::ColumnParallel, LoraSide::A)
+        | (LoraTpShardLayout::FlatQkvColumnParallel { .. }, LoraSide::A)
+        | (LoraTpShardLayout::RowParallel, LoraSide::B) => {
+            if local_shape[rank_axis] != global_lora_rank {
+                bail!(
+                    "replicated checkpoint tensor {file}:{tensor_name} has rank {} on axis {rank_axis}, expected {global_lora_rank}",
+                    local_shape[rank_axis]
+                );
+            }
+            (rank_axis, true)
+        }
+        (LoraTpShardLayout::ColumnParallel, LoraSide::B) => {
+            if local_shape[rank_axis] != global_lora_rank {
+                bail!(
+                    "column-parallel checkpoint tensor {file}:{tensor_name} has rank {} on axis {rank_axis}, expected {global_lora_rank}",
+                    local_shape[rank_axis]
+                );
+            }
+            let axis = local_shape.len() - 2;
+            global_shape[axis] *= tp_size;
+            global_offset[axis] = tp_rank * local_shape[axis];
+            (axis, false)
+        }
+        (
+            LoraTpShardLayout::FlatQkvColumnParallel {
+                q_rows,
+                k_rows,
+                v_rows,
+            },
+            LoraSide::B,
+        ) => {
+            if local_shape[rank_axis] != global_lora_rank {
+                bail!(
+                    "flat-QKV column-parallel checkpoint tensor {file}:{tensor_name} has rank {} on axis {rank_axis}, expected {global_lora_rank}",
+                    local_shape[rank_axis]
+                );
+            }
+            if q_rows <= 0
+                || k_rows <= 0
+                || v_rows <= 0
+                || q_rows % tp_size != 0
+                || k_rows % tp_size != 0
+                || v_rows % tp_size != 0
+            {
+                bail!(
+                    "flat-QKV global row segments [{q_rows}, {k_rows}, {v_rows}] must be positive and divisible by TP size {tp_size}"
+                );
+            }
+            let axis = local_shape.len() - 2;
+            let local_q = q_rows / tp_size;
+            let local_k = k_rows / tp_size;
+            let local_v = v_rows / tp_size;
+            if local_shape[axis] != local_q + local_k + local_v {
+                bail!(
+                    "flat-QKV checkpoint tensor {file}:{tensor_name} has {} local rows, expected {} from Q/K/V segments",
+                    local_shape[axis],
+                    local_q + local_k + local_v
+                );
+            }
+            global_shape[axis] = q_rows + k_rows + v_rows;
+            (axis, false)
+        }
+        (LoraTpShardLayout::RowParallel, LoraSide::A) => {
+            if local_shape[rank_axis] != global_lora_rank {
+                bail!(
+                    "row-parallel checkpoint tensor {file}:{tensor_name} has rank {} on axis {rank_axis}, expected {global_lora_rank}",
+                    local_shape[rank_axis]
+                );
+            }
+            let axis = local_shape.len() - 1;
+            global_shape[axis] *= tp_size;
+            global_offset[axis] = tp_rank * local_shape[axis];
+            (axis, false)
+        }
+        (LoraTpShardLayout::RoutedExpertFusedGateUp, LoraSide::A)
+        | (LoraTpShardLayout::RoutedExpertDown, LoraSide::B) => {
+            if local_shape.len() != 3 || local_shape[rank_axis] != global_lora_rank {
+                bail!(
+                    "routed expert checkpoint tensor {file}:{tensor_name} must be rank 3 with LoRA rank {global_lora_rank}"
+                );
+            }
+            (rank_axis, true)
+        }
+        (LoraTpShardLayout::RoutedExpertFusedGateUp, LoraSide::B) => {
+            if local_shape.len() != 3 || local_shape[rank_axis] != global_lora_rank {
+                bail!(
+                    "routed fused gate/up checkpoint tensor {file}:{tensor_name} must be rank 3 with LoRA rank {global_lora_rank}"
+                );
+            }
+            let axis = local_shape.len() - 2;
+            if local_shape[axis] <= 0 || local_shape[axis] % 2 != 0 {
+                bail!(
+                    "routed fused gate/up checkpoint tensor {file}:{tensor_name} must have an even positive local projection size"
+                );
+            }
+            global_shape[axis] *= tp_size;
+            (axis, false)
+        }
+        (LoraTpShardLayout::RoutedExpertDown, LoraSide::A) => {
+            if local_shape.len() != 3 || local_shape[rank_axis] != global_lora_rank {
+                bail!(
+                    "routed expert down checkpoint tensor {file}:{tensor_name} must be rank 3 with LoRA rank {global_lora_rank}"
+                );
+            }
+            let axis = local_shape.len() - 1;
+            global_shape[axis] *= tp_size;
+            global_offset[axis] = tp_rank * local_shape[axis];
+            (axis, false)
+        }
+    };
+    let segments = match (layout, side) {
+        (
+            LoraTpShardLayout::FlatQkvColumnParallel {
+                q_rows,
+                k_rows,
+                v_rows,
+            },
+            LoraSide::B,
+        ) => {
+            let local_q = q_rows / tp_size;
+            let local_k = k_rows / tp_size;
+            let local_v = v_rows / tp_size;
+            vec![
+                TensorShardSegmentManifest {
+                    local_offset: 0,
+                    global_offset: tp_rank * local_q,
+                    length: local_q,
+                },
+                TensorShardSegmentManifest {
+                    local_offset: local_q,
+                    global_offset: q_rows + tp_rank * local_k,
+                    length: local_k,
+                },
+                TensorShardSegmentManifest {
+                    local_offset: local_q + local_k,
+                    global_offset: q_rows + k_rows + tp_rank * local_v,
+                    length: local_v,
+                },
+            ]
+        }
+        (LoraTpShardLayout::RoutedExpertFusedGateUp, LoraSide::B) => {
+            let axis = local_shape.len() - 2;
+            let local_half = local_shape[axis] / 2;
+            let global_half = local_half * tp_size;
+            vec![
+                TensorShardSegmentManifest {
+                    local_offset: 0,
+                    global_offset: tp_rank * local_half,
+                    length: local_half,
+                },
+                TensorShardSegmentManifest {
+                    local_offset: local_half,
+                    global_offset: global_half + tp_rank * local_half,
+                    length: local_half,
+                },
+            ]
+        }
+        _ => Vec::new(),
+    };
+    let routed_expert = matches!(
+        layout,
+        LoraTpShardLayout::RoutedExpertFusedGateUp | LoraTpShardLayout::RoutedExpertDown
+    );
+    let ep_size = i64::try_from(parallel.expert_model_parallel_size)
+        .context("EP size exceeds checkpoint tensor shape range")?;
+    let ep_rank = i64::try_from(parallel.coordinates.expert)
+        .context("EP rank exceeds checkpoint tensor shape range")?;
+    let mut placements = Vec::new();
+    let mut replicated_axes = Vec::new();
+    if routed_expert && ep_size > 1 {
+        global_shape[0] *= ep_size;
+        global_offset[0] = ep_rank * local_shape[0];
+        placements.push(TensorShardPlacementManifest {
+            parallel_axis: ParallelAxis::Expert,
+            tensor_axis: 0,
+            global_size: global_shape[0],
+            local_size: local_shape[0],
+            global_offset: global_offset[0],
+            segments: Vec::new(),
+        });
+    } else if parallel.expert_model_parallel_size > 1 {
+        replicated_axes.push(ParallelAxis::Expert);
+    }
+    if parallel.tensor_model_parallel_size > 1 {
+        if replicated {
+            replicated_axes.push(ParallelAxis::Tensor);
+        } else {
+            placements.push(TensorShardPlacementManifest {
+                parallel_axis: ParallelAxis::Tensor,
+                tensor_axis: partition_axis,
+                global_size: global_shape[partition_axis],
+                local_size: local_shape[partition_axis],
+                global_offset: global_offset[partition_axis],
+                segments: segments.clone(),
+            });
+        }
+    }
+    if parallel.data_parallel_size > 1 {
+        replicated_axes.push(ParallelAxis::Data);
+    }
+    Ok(TensorShardManifest {
+        file: file.to_string(),
+        tensor_name,
+        state: state.to_string(),
+        adapter_id,
+        global_lora_rank,
+        global_shape,
+        local_shape,
+        partition_axis,
+        layout,
+        replicated,
+        global_offset,
+        segments,
+        placements,
+        replicated_axes,
+        replica_identity: if replicated {
+            "tp-replicated".to_string()
+        } else {
+            parallel.replica_identity()
+        },
+    })
+}
+
+fn validate_saved_shards(
+    saved: &[TensorShardManifest],
+    expected: &[TensorShardManifest],
+) -> Result<()> {
+    let keyed = |shards: &[TensorShardManifest]| -> Result<BTreeMap<_, _>> {
+        let mut by_key = BTreeMap::new();
+        for shard in shards {
+            let key = (
+                shard.file.clone(),
+                shard.tensor_name.clone(),
+                shard.state.clone(),
+            );
+            if by_key.insert(key, shard.clone()).is_some() {
+                bail!(
+                    "duplicate tensor shard metadata for {}:{} ({})",
+                    shard.file,
+                    shard.tensor_name,
+                    shard.state
+                );
+            }
+        }
+        Ok(by_key)
+    };
+    let saved = keyed(saved)?;
+    let expected = keyed(expected)?;
+    if saved.len() != expected.len() {
+        bail!(
+            "tensor shard metadata count mismatch: saved={}, expected={}",
+            saved.len(),
+            expected.len()
+        );
+    }
+    for (key, expected_shard) in expected {
+        let saved_shard = saved
+            .get(&key)
+            .with_context(|| format!("missing tensor shard metadata for {}:{}", key.0, key.1))?;
+        if saved_shard != &expected_shard {
+            bail!(
+                "tensor shard metadata mismatch for {}:{}: saved={saved_shard:?}, expected={expected_shard:?}",
+                key.0,
+                key.1
+            );
+        }
+    }
+    Ok(())
+}
+
+fn rank_checkpoint_dir(root: &Path, global_rank: usize) -> PathBuf {
+    root.join(format!("rank-{global_rank:05}"))
+}
+
+fn env_usize(names: &[&str], default: usize) -> Result<usize> {
+    Ok(env_usize_optional(names)?.unwrap_or(default))
+}
+
+fn env_usize_optional(names: &[&str]) -> Result<Option<usize>> {
+    let Some((name, value)) = names
+        .iter()
+        .find_map(|name| env::var(name).ok().map(|value| (*name, value)))
+    else {
+        return Ok(None);
+    };
+    let value = value
+        .parse::<usize>()
+        .with_context(|| format!("{name} must be a non-negative integer"))?;
+    if value == 0 && name != "RANK" {
+        bail!("{name} must be positive");
+    }
+    Ok(Some(value))
+}
+
+fn named_tensors(
+    a: &[Tensor],
+    b: &[Tensor],
+    a_prefix: &str,
+    b_prefix: &str,
+) -> Vec<(String, Tensor)> {
+    let mut named: Vec<(String, Tensor)> = Vec::new();
+    for (i, t) in a.iter().enumerate() {
+        named.push((
+            format!("{a_prefix}{i}"),
+            t.to_kind(tch::Kind::Float).to_device(tch::Device::Cpu),
+        ));
+    }
+    for (i, t) in b.iter().enumerate() {
+        named.push((
+            format!("{b_prefix}{i}"),
+            t.to_kind(tch::Kind::Float).to_device(tch::Device::Cpu),
+        ));
+    }
+
+    named
+}
+
+fn save_named_tensors(path: &Path, named: Vec<(String, Tensor)>) -> Result<()> {
+    Tensor::write_safetensors(&named, path).with_context(|| format!("write {}", path.display()))?;
+    Ok(())
+}
+
+fn read_named_tensors(path: &Path) -> Result<std::collections::BTreeMap<String, Tensor>> {
+    let named =
+        Tensor::read_safetensors(path).with_context(|| format!("read {}", path.display()))?;
+    let mut by_name = std::collections::BTreeMap::new();
+    for (name, tensor) in named {
+        by_name.insert(name, tensor);
+    }
+    Ok(by_name)
+}
+
+fn collect_side(
+    tensors: &std::collections::BTreeMap<String, Tensor>,
+    prefix: &str,
+) -> Result<Vec<Tensor>> {
+    let mut indices = tensors
+        .keys()
+        .filter_map(|name| name.strip_prefix(prefix)?.parse::<usize>().ok())
+        .collect::<Vec<_>>();
+    indices.sort_unstable();
+    let mut result = Vec::with_capacity(indices.len());
+    for (expected, index) in indices.into_iter().enumerate() {
+        if index != expected {
+            anyhow::bail!("checkpoint tensor indices for {prefix} are not contiguous");
+        }
+        result.push(
+            tensors
+                .get(&format!("{prefix}{index}"))
+                .expect("index collected from map")
+                .shallow_clone(),
+        );
+    }
+    Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[allow(clippy::too_many_arguments)]
+    fn save_checkpoint_with_dynamic_for_topology(
+        dir: &Path,
+        step: u64,
+        loss: f64,
+        model_path: &str,
+        lora_rank: i64,
+        lora_alpha: f64,
+        lora_a: &[Tensor],
+        lora_b: &[Tensor],
+        adam_m: &[Tensor],
+        adam_v: &[Tensor],
+        dynamic_adapters: &[DynamicAdapterCheckpoint],
+        fixed_shard_layouts: &[LoraTpShardLayout],
+        fixed_slot_identities: &[LoraSlotIdentity],
+        parallel: &ParallelCheckpointManifest,
+    ) -> Result<()> {
+        super::save_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+            dir,
+            step,
+            if lora_a.is_empty() { 0 } else { step },
+            loss,
+            model_path,
+            lora_rank,
+            lora_alpha,
+            lora_a,
+            lora_b,
+            adam_m,
+            adam_v,
+            dynamic_adapters,
+            fixed_shard_layouts,
+            fixed_slot_identities,
+            parallel,
+            Some("checkpoint-test-generation"),
+        )
+    }
+
+    #[test]
+    fn partial_rank_manifest_is_not_ready() {
+        let root = tempfile::tempdir().unwrap();
+        let rank_dir = root.path().join("rank-00000");
+        std::fs::create_dir(&rank_dir).unwrap();
+        let partial = rank_dir.join(".manifest.json.partial");
+        std::fs::write(&partial, b"{\"format\":").unwrap();
+
+        assert!(!rank_manifests_ready(root.path(), 1));
+
+        std::fs::rename(partial, rank_dir.join("manifest.json")).unwrap();
+        std::fs::write(rank_dir.join(RANK_RECEIPT_FILE), b"{}").unwrap();
+        assert!(rank_manifests_ready(root.path(), 1));
+    }
+
+    #[test]
+    fn distributed_checkpoint_requires_nonempty_generation() {
+        let root = tempfile::tempdir().unwrap();
+        let parallel = ep_topology(0, 2);
+
+        let error = super::save_checkpoint_with_dynamic_and_fixed_step_for_topology(
+            root.path(),
+            0,
+            0,
+            0.0,
+            "Qwen/test",
+            2,
+            4.0,
+            &[],
+            &[],
+            &[],
+            &[],
+            &[],
+            &[],
+            &[],
+            &parallel,
+        )
+        .unwrap_err();
+        assert!(error.to_string().contains("explicit unique generation"));
+
+        for generation in [None, Some("")] {
+            let error = super::save_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+                root.path(),
+                0,
+                0,
+                0.0,
+                "Qwen/test",
+                2,
+                4.0,
+                &[],
+                &[],
+                &[],
+                &[],
+                &[],
+                &[],
+                &[],
+                &parallel,
+                generation,
+            )
+            .unwrap_err();
+            assert!(error
+                .to_string()
+                .contains("distributed checkpoint generation must be non-empty"));
+        }
+
+        let too_long = "x".repeat(257);
+        for generation in ["bad/generation", too_long.as_str()] {
+            let error = super::save_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+                root.path(),
+                0,
+                0,
+                0.0,
+                "Qwen/test",
+                2,
+                4.0,
+                &[],
+                &[],
+                &[],
+                &[],
+                &[],
+                &[],
+                &[],
+                &parallel,
+                Some(generation),
+            )
+            .unwrap_err();
+            assert!(error.to_string().contains("path-safe ASCII"));
+        }
+
+        super::save_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+            root.path(),
+            0,
+            0,
+            0.0,
+            "Qwen/test",
+            2,
+            4.0,
+            &[],
+            &[],
+            &[],
+            &[],
+            &[],
+            &[],
+            &[],
+            &parallel,
+            Some("save-transaction-1"),
+        )
+        .unwrap();
+        let error = super::save_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+            root.path(),
+            0,
+            0,
+            0.0,
+            "Qwen/test",
+            2,
+            4.0,
+            &[],
+            &[],
+            &[],
+            &[],
+            &[],
+            &[],
+            &[],
+            &parallel,
+            Some("save-transaction-1"),
+        )
+        .unwrap_err();
+        assert!(error.to_string().contains("has already been used"));
+    }
+
+    #[test]
+    fn distributed_export_publishes_one_peft_artifact() {
+        let root = tempfile::tempdir().unwrap();
+        let model_dir = root.path().join("model");
+        std::fs::create_dir(&model_dir).unwrap();
+        std::fs::write(
+            model_dir.join("config.json"),
+            r#"{
+                "model_type": "qwen3_5_text",
+                "num_hidden_layers": 1,
+                "hidden_size": 4,
+                "vocab_size": 16,
+                "num_attention_heads": 1,
+                "num_key_value_heads": 1,
+                "head_dim": 4,
+                "intermediate_size": 8
+            }"#,
+        )
+        .unwrap();
+        let final_path = root.path().join("adapter");
+        for rank in 0..2 {
+            let stale_rank = root.path().join(format!(
+                "adapter.rustrain-shards-old-generation/rank-{rank:05}"
+            ));
+            std::fs::create_dir_all(&stale_rank).unwrap();
+            std::fs::write(stale_rank.join("manifest.json"), "stale").unwrap();
+        }
+        let model_path = model_dir.to_string_lossy().into_owned();
+        let results = std::thread::scope(|scope| {
+            let mut handles = Vec::new();
+            for global_rank in 0..2 {
+                let final_path = final_path.clone();
+                let model_path = model_path.clone();
+                handles.push(scope.spawn(move || {
+                    let parallel = tp_topology(global_rank, 2);
+                    let a = Tensor::full([2, 4], 3.0, (tch::Kind::Float, tch::Device::Cpu));
+                    let b = Tensor::full(
+                        [2, 2],
+                        global_rank as f64 + 1.0,
+                        (tch::Kind::Float, tch::Device::Cpu),
+                    );
+                    let adam_m = vec![a.zeros_like(), b.zeros_like()];
+                    let adam_v = vec![a.ones_like(), b.ones_like()];
+                    export_distributed_adapter_checkpoint(
+                        &final_path,
+                        "test-generation",
+                        &parallel,
+                        None,
+                        |staging| {
+                            save_checkpoint_with_dynamic_for_topology(
+                                staging,
+                                7,
+                                0.25,
+                                &model_path,
+                                2,
+                                4.0,
+                                &[a],
+                                &[b],
+                                &adam_m,
+                                &adam_v,
+                                &[],
+                                &[LoraTpShardLayout::ColumnParallel],
+                                &[LoraSlotIdentity {
+                                    index: 0,
+                                    layer: 0,
+                                    module: "q_proj".to_string(),
+                                }],
+                                &parallel,
+                            )
+                        },
+                    )
+                }));
+            }
+            handles
+                .into_iter()
+                .map(|handle| handle.join().unwrap().unwrap())
+                .collect::<Vec<_>>()
+        });
+
+        assert_eq!(results, vec![2, 2]);
+        let artifact = Qwen36AdapterArtifact::load(&final_path).unwrap();
+        assert_eq!(artifact.tensors.len(), 2);
+        let a = artifact
+            .tensors
+            .get("base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight")
+            .unwrap();
+        let b = artifact
+            .tensors
+            .get("base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight")
+            .unwrap();
+        assert_eq!(a.size(), [2, 4]);
+        assert_eq!(b.size(), [4, 2]);
+        assert_eq!(b.double_value(&[0, 0]), 1.0);
+        assert_eq!(b.double_value(&[3, 0]), 2.0);
+        assert!(final_path.join(".rustrain-export-completed.json").is_file());
+    }
+
+    #[test]
+    fn distributed_export_rejects_reused_attempt_id() {
+        let root = tempfile::tempdir().unwrap();
+        std::fs::create_dir(root.path().join("adapter.rustrain-shards-reused")).unwrap();
+        let error = export_distributed_adapter_checkpoint(
+            &root.path().join("adapter"),
+            "reused",
+            &tp_topology(0, 2),
+            None,
+            |_| panic!("reused attempt must fail before writing rank state"),
+        )
+        .unwrap_err();
+        assert!(error.to_string().contains("attempt IDs cannot be reused"));
+        assert!(!root
+            .path()
+            .join("adapter.rustrain-shards-reused/errors")
+            .exists());
+    }
+
+    #[test]
+    fn distributed_export_rejects_file_paths() {
+        let root = tempfile::tempdir().unwrap();
+        let error = export_distributed_adapter_checkpoint(
+            &root.path().join("adapter.safetensors"),
+            "test-generation",
+            &tp_topology(0, 2),
+            None,
+            |_| panic!("invalid export path must fail before writing rank state"),
+        )
+        .unwrap_err();
+        assert!(error.to_string().contains("must be a directory"));
+    }
+
+    #[test]
+    fn checkpoint_adapter_and_optimizer_roundtrip() {
+        let dir = tempfile::tempdir().unwrap();
+        let a = Tensor::arange(6, (tch::Kind::Float, tch::Device::Cpu)).reshape([2, 3]);
+        let b = Tensor::arange(8, (tch::Kind::Float, tch::Device::Cpu)).reshape([4, 2]);
+        let m = Tensor::ones([2, 3], (tch::Kind::Float, tch::Device::Cpu));
+        let v = Tensor::full([4, 2], 2.0, (tch::Kind::Float, tch::Device::Cpu));
+        save_checkpoint(
+            dir.path(),
+            7,
+            1.25,
+            "Qwen/test",
+            2,
+            4.0,
+            &[a.shallow_clone()],
+            &[b.shallow_clone()],
+            &[m.shallow_clone()],
+            &[v.shallow_clone()],
+        )
+        .unwrap();
+        let loaded = load_checkpoint(dir.path()).unwrap();
+        assert!(loaded.manifest.parallel.is_none());
+        assert!(loaded.manifest.tensor_shards.is_empty());
+        assert_eq!(loaded.manifest.lora_alpha, 4.0);
+        assert_eq!(loaded.manifest.step, 7);
+        assert_eq!(loaded.manifest.effective_fixed_optimizer_step(), 7);
+        assert_eq!(loaded.lora_a[0].size(), [2, 3]);
+        assert_eq!(loaded.lora_b[0].size(), [4, 2]);
+        assert!(loaded.lora_a[0].allclose(&a, 1e-6, 1e-6, false));
+        assert!(loaded.lora_b[0].allclose(&b, 1e-6, 1e-6, false));
+        assert!(loaded.adam_m[0].allclose(&m, 1e-6, 1e-6, false));
+        assert!(loaded.adam_v[0].allclose(&v, 1e-6, 1e-6, false));
+
+        let manifest_path = dir.path().join("manifest.json");
+        let mut legacy: serde_json::Value =
+            serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap();
+        legacy
+            .as_object_mut()
+            .unwrap()
+            .remove("fixed_optimizer_step");
+        std::fs::write(
+            &manifest_path,
+            serde_json::to_string_pretty(&legacy).unwrap(),
+        )
+        .unwrap();
+        let legacy = load_checkpoint(dir.path()).unwrap();
+        assert_eq!(legacy.manifest.fixed_optimizer_step, None);
+        assert_eq!(legacy.manifest.effective_fixed_optimizer_step(), 7);
+
+        let split_clock_dir = tempfile::tempdir().unwrap();
+        let single_rank = ParallelCheckpointManifest::new(1, 0, 1, 1, 1, 1, 1).unwrap();
+        save_checkpoint_with_dynamic_and_fixed_step_for_topology(
+            split_clock_dir.path(),
+            9,
+            3,
+            1.0,
+            "Qwen/test",
+            2,
+            4.0,
+            &[a],
+            &[b],
+            &[m],
+            &[v],
+            &[],
+            &[],
+            &[],
+            &single_rank,
+        )
+        .unwrap();
+        let split_clock = load_checkpoint(split_clock_dir.path()).unwrap();
+        assert_eq!(split_clock.manifest.step, 9);
+        assert_eq!(split_clock.manifest.effective_fixed_optimizer_step(), 3);
+    }
+
+    #[test]
+    fn dynamic_checkpoint_roundtrip_preserves_metadata_and_state() {
+        let dir = tempfile::tempdir().unwrap();
+        let dynamic = DynamicAdapterCheckpoint {
+            manifest: DynamicAdapterManifest {
+                id: 7,
+                rank: 3,
+                alpha: 6.0,
+                optimizer_step: 19,
+                optimizer_lr: Some(2.5e-4),
+                optimizer_beta1: Some(0.8),
+                optimizer_beta2: Some(0.95),
+                optimizer_eps: Some(1e-6),
+                target_layers: vec![1, 3],
+                target_modules: vec!["q_proj".into(), "down_proj".into()],
+                shard_layouts: Vec::new(),
+                slot_identities: Vec::new(),
+                parameter_count: 2,
+                optimizer_count: 4,
+            },
+            lora_a: (0..2)
+                .map(|_| Tensor::ones([3, 8], (tch::Kind::Float, tch::Device::Cpu)))
+                .collect(),
+            lora_b: (0..2)
+                .map(|_| Tensor::full([16, 3], 2.0, (tch::Kind::Float, tch::Device::Cpu)))
+                .collect(),
+            adam_m: (0..4)
+                .map(|_| Tensor::full([3, 8], 3.0, (tch::Kind::Float, tch::Device::Cpu)))
+                .collect(),
+            adam_v: (0..4)
+                .map(|_| Tensor::full([3, 8], 4.0, (tch::Kind::Float, tch::Device::Cpu)))
+                .collect(),
+        };
+        save_checkpoint_with_dynamic(
+            dir.path(),
+            11,
+            0.5,
+            "Qwen/test",
+            2,
+            4.0,
+            &[],
+            &[],
+            &[],
+            &[],
+            &[dynamic],
+        )
+        .unwrap();
+        let loaded = load_checkpoint(dir.path()).unwrap();
+        assert_eq!(loaded.manifest.format, "rustrain-checkpoint-v2");
+        assert!(loaded.manifest.parallel.is_none());
+        assert!(loaded.manifest.tensor_shards.is_empty());
+        assert_eq!(loaded.dynamic_adapters.len(), 1);
+        let loaded_dynamic = &loaded.dynamic_adapters[0];
+        assert_eq!(loaded_dynamic.manifest.id, 7);
+        assert_eq!(loaded_dynamic.manifest.rank, 3);
+        assert_eq!(loaded_dynamic.manifest.optimizer_step, 19);
+        assert_eq!(loaded_dynamic.manifest.optimizer_lr, Some(2.5e-4));
+        assert_eq!(loaded_dynamic.manifest.optimizer_beta1, Some(0.8));
+        assert_eq!(loaded_dynamic.manifest.optimizer_beta2, Some(0.95));
+        assert_eq!(loaded_dynamic.manifest.optimizer_eps, Some(1e-6));
+        assert_eq!(loaded_dynamic.manifest.target_layers, vec![1, 3]);
+        assert_eq!(loaded_dynamic.lora_a.len(), 2);
+        assert_eq!(loaded_dynamic.adam_m.len(), 4);
+        assert!(loaded_dynamic.adam_m[0].allclose(
+            &Tensor::full([3, 8], 3.0, (tch::Kind::Float, tch::Device::Cpu)),
+            1e-6,
+            1e-6,
+            false
+        ));
+    }
+
+    #[test]
+    fn old_dynamic_manifest_defaults_optimizer_step_to_zero() {
+        let json = r#"{
+            "id": 7,
+            "rank": 3,
+            "alpha": 6.0,
+            "target_layers": [1],
+            "target_modules": ["q_proj"],
+            "parameter_count": 1,
+            "optimizer_count": 2
+        }"#;
+        let manifest: DynamicAdapterManifest = serde_json::from_str(json).unwrap();
+        assert_eq!(manifest.optimizer_step, 0);
+        assert_eq!(manifest.optimizer_lr, None);
+        assert!(manifest.shard_layouts.is_empty());
+    }
+
+    #[test]
+    fn optimizer_clocks_require_matching_adam_state() {
+        let mut manifest = resume_validation_manifest("rustrain-checkpoint-v2");
+        manifest.fixed_optimizer_step = Some(2);
+        let error = validate_optimizer_clock_state(&manifest, &[], &[], &[]).unwrap_err();
+        assert!(error.to_string().contains("fixed optimizer step 2"));
+
+        manifest.fixed_optimizer_step = Some(0);
+        let dynamic = DynamicAdapterCheckpoint {
+            manifest: DynamicAdapterManifest {
+                id: 7,
+                rank: 4,
+                alpha: 8.0,
+                optimizer_step: 3,
+                optimizer_lr: None,
+                optimizer_beta1: None,
+                optimizer_beta2: None,
+                optimizer_eps: None,
+                target_layers: vec![0],
+                target_modules: vec!["q_proj".to_string()],
+                shard_layouts: Vec::new(),
+                slot_identities: Vec::new(),
+                parameter_count: 0,
+                optimizer_count: 0,
+            },
+            lora_a: Vec::new(),
+            lora_b: Vec::new(),
+            adam_m: Vec::new(),
+            adam_v: Vec::new(),
+        };
+        let error = validate_optimizer_clock_state(&manifest, &[], &[], &[dynamic]).unwrap_err();
+        assert!(error
+            .to_string()
+            .contains("dynamic adapter 7 optimizer step 3"));
+    }
+
+    fn tp_topology(global_rank: usize, tp_size: usize) -> ParallelCheckpointManifest {
+        ParallelCheckpointManifest::new(tp_size, global_rank, tp_size, 1, 1, 1, 1).unwrap()
+    }
+
+    fn pp_topology(global_rank: usize, pp_size: usize) -> ParallelCheckpointManifest {
+        ParallelCheckpointManifest::new(pp_size, global_rank, 1, pp_size, 1, 1, 1).unwrap()
+    }
+
+    fn pp_stage_metadata(pipeline_rank: usize) -> StageUnionCheckpointMetadata {
+        StageUnionCheckpointMetadata {
+            pipeline_stage: PipelineStageCheckpointManifest {
+                pipeline_rank,
+                pipeline_size: 2,
+                global_num_layers: 2,
+                layer_start: pipeline_rank,
+                layer_end: pipeline_rank + 1,
+            },
+            fixed_target_layers: vec![0, 1],
+            fixed_target_modules: vec!["q_proj".to_string()],
+        }
+    }
+
+    #[test]
+    fn stage_union_checkpoint_round_trips_stage_local_global_identities() {
+        let root = tempfile::tempdir().unwrap();
+        for pipeline_rank in 0..2 {
+            let topology = pp_topology(pipeline_rank, 2);
+            let value = pipeline_rank as f64 + 1.0;
+            let a = Tensor::full([2, 4], value, (tch::Kind::Float, tch::Device::Cpu));
+            let b = Tensor::full([4, 2], value + 10.0, (tch::Kind::Float, tch::Device::Cpu));
+            let adam_m = vec![Tensor::zeros_like(&a), Tensor::zeros_like(&b)];
+            let adam_v = vec![Tensor::ones_like(&a), Tensor::ones_like(&b)];
+            let identity = LoraSlotIdentity {
+                index: pipeline_rank * 4,
+                layer: pipeline_rank,
+                module: "q_proj".to_string(),
+            };
+            save_stage_union_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+                root.path(),
+                7,
+                3,
+                0.25,
+                "Qwen/test",
+                2,
+                4.0,
+                &[a],
+                &[b],
+                &adam_m,
+                &adam_v,
+                &[],
+                &[LoraTpShardLayout::LatentRank],
+                &[identity],
+                &topology,
+                "pp-stage-union-test",
+                &pp_stage_metadata(pipeline_rank),
+            )
+            .unwrap();
+        }
+
+        for pipeline_rank in 0..2 {
+            let topology = pp_topology(pipeline_rank, 2);
+            let loaded = load_checkpoint_for_topology(root.path(), &topology).unwrap();
+            assert_eq!(loaded.manifest.format, STAGE_UNION_CHECKPOINT_FORMAT);
+            assert_eq!(loaded.manifest.fixed_slot_identities.len(), 1);
+            assert_eq!(
+                loaded.manifest.fixed_slot_identities[0],
+                LoraSlotIdentity {
+                    index: pipeline_rank * 4,
+                    layer: pipeline_rank,
+                    module: "q_proj".to_string(),
+                }
+            );
+            assert_eq!(
+                loaded.lora_a[0].double_value(&[0, 0]),
+                pipeline_rank as f64 + 1.0
+            );
+        }
+    }
+
+    #[test]
+    fn stage_union_legacy_dynamic_digest_keeps_receipt_compatibility() {
+        let mut manifest = resume_validation_manifest(STAGE_UNION_CHECKPOINT_FORMAT);
+        manifest.checkpoint_generation = Some("legacy-stage-union".to_string());
+        manifest.fixed_target_layers = vec![0];
+        manifest.fixed_target_modules = vec!["q_proj".to_string()];
+        manifest.dynamic_adapters = vec![DynamicAdapterManifest {
+            id: 7,
+            rank: 4,
+            alpha: 8.0,
+            optimizer_step: 3,
+            optimizer_lr: Some(5e-4),
+            optimizer_beta1: None,
+            optimizer_beta2: None,
+            optimizer_eps: None,
+            target_layers: vec![0],
+            target_modules: vec!["q_proj".to_string()],
+            shard_layouts: Vec::new(),
+            slot_identities: Vec::new(),
+            parameter_count: 0,
+            optimizer_count: 0,
+        }];
+        let expected_dynamic = manifest
+            .dynamic_adapters
+            .iter()
+            .map(|adapter| LegacyStageUnionDynamicAdapterIdentity {
+                id: adapter.id,
+                rank: adapter.rank,
+                alpha_bits: adapter.alpha.to_bits(),
+                optimizer_step: adapter.optimizer_step,
+                optimizer_lr_bits: adapter.optimizer_lr.map(f64::to_bits),
+                target_layers: &adapter.target_layers,
+                target_modules: &adapter.target_modules,
+            })
+            .collect();
+        let expected = stable_bytes_digest(
+            &serde_json::to_vec(&LegacyStageUnionGenerationIdentity {
+                checkpoint_generation: &manifest.checkpoint_generation,
+                step: manifest.step,
+                fixed_optimizer_step: manifest.effective_fixed_optimizer_step(),
+                model_path: &manifest.model_path,
+                lora_rank: manifest.lora_rank,
+                lora_alpha_bits: manifest.lora_alpha.to_bits(),
+                files: &manifest.files,
+                fixed_target_layers: &manifest.fixed_target_layers,
+                fixed_target_modules: &manifest.fixed_target_modules,
+                dynamic_adapters: expected_dynamic,
+            })
+            .unwrap(),
+        );
+        assert_eq!(checkpoint_generation_digest(&manifest).unwrap(), expected);
+    }
+
+    #[test]
+    fn stage_union_save_rejects_stage_topology_mismatch_before_writing() {
+        let root = tempfile::tempdir().unwrap();
+        let topology = pp_topology(0, 2);
+        let error =
+            save_stage_union_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+                root.path(),
+                0,
+                0,
+                0.0,
+                "Qwen/test",
+                2,
+                4.0,
+                &[],
+                &[],
+                &[],
+                &[],
+                &[],
+                &[],
+                &[],
+                &topology,
+                "pp-stage-union-invalid",
+                &pp_stage_metadata(1),
+            )
+            .unwrap_err();
+        assert!(error.to_string().contains("does not match topology stage"));
+        assert!(!root.path().join("rank-00000").exists());
+    }
+
+    fn ep_topology(global_rank: usize, ep_size: usize) -> ParallelCheckpointManifest {
+        ParallelCheckpointManifest::new(ep_size, global_rank, 1, 1, 1, ep_size, 1).unwrap()
+    }
+
+    fn tp_ep_topology(global_rank: usize) -> ParallelCheckpointManifest {
+        ParallelCheckpointManifest::new(4, global_rank, 2, 1, 1, 2, 1).unwrap()
+    }
+
+    fn tp_ep_dp_topology(global_rank: usize) -> ParallelCheckpointManifest {
+        ParallelCheckpointManifest::new(8, global_rank, 2, 1, 2, 2, 1).unwrap()
+    }
+
+    fn save_tp_ep_dp_checkpoint(dir: &Path) {
+        let identity = LoraSlotIdentity {
+            index: 0,
+            layer: 0,
+            module: "experts_gate_up_proj".to_string(),
+        };
+        for global_rank in 0..8 {
+            let topology = tp_ep_dp_topology(global_rank);
+            let value = (topology.coordinates.expert * 10 + topology.coordinates.tensor) as f64;
+            let a = Tensor::full([2, 4, 2], value + 1.0, (tch::Kind::Float, tch::Device::Cpu));
+            let b = Tensor::full([2, 4, 4], value + 2.0, (tch::Kind::Float, tch::Device::Cpu));
+            let adam_m = vec![
+                Tensor::full(a.size(), value + 3.0, (tch::Kind::Float, tch::Device::Cpu)),
+                Tensor::full(b.size(), value + 4.0, (tch::Kind::Float, tch::Device::Cpu)),
+            ];
+            let adam_v = vec![
+                Tensor::full(a.size(), value + 5.0, (tch::Kind::Float, tch::Device::Cpu)),
+                Tensor::full(b.size(), value + 6.0, (tch::Kind::Float, tch::Device::Cpu)),
+            ];
+            let dynamic = DynamicAdapterCheckpoint {
+                manifest: DynamicAdapterManifest {
+                    id: 17,
+                    rank: 4,
+                    alpha: 8.0,
+                    optimizer_step: 29,
+                    optimizer_lr: Some(1e-3),
+                    optimizer_beta1: Some(0.85),
+                    optimizer_beta2: Some(0.97),
+                    optimizer_eps: Some(1e-7),
+                    target_layers: vec![0],
+                    target_modules: vec!["experts_gate_up_proj".to_string()],
+                    shard_layouts: vec![LoraTpShardLayout::RoutedExpertFusedGateUp],
+                    slot_identities: vec![identity.clone()],
+                    parameter_count: 1,
+                    optimizer_count: 2,
+                },
+                lora_a: vec![Tensor::full(
+                    [2, 4, 2],
+                    value + 11.0,
+                    (tch::Kind::Float, tch::Device::Cpu),
+                )],
+                lora_b: vec![Tensor::full(
+                    [2, 4, 4],
+                    value + 12.0,
+                    (tch::Kind::Float, tch::Device::Cpu),
+                )],
+                adam_m: vec![
+                    Tensor::full(
+                        [2, 4, 2],
+                        value + 13.0,
+                        (tch::Kind::Float, tch::Device::Cpu),
+                    ),
+                    Tensor::full(
+                        [2, 4, 4],
+                        value + 14.0,
+                        (tch::Kind::Float, tch::Device::Cpu),
+                    ),
+                ],
+                adam_v: vec![
+                    Tensor::full(
+                        [2, 4, 2],
+                        value + 15.0,
+                        (tch::Kind::Float, tch::Device::Cpu),
+                    ),
+                    Tensor::full(
+                        [2, 4, 4],
+                        value + 16.0,
+                        (tch::Kind::Float, tch::Device::Cpu),
+                    ),
+                ],
+            };
+            save_checkpoint_with_dynamic_for_topology(
+                dir,
+                31,
+                0.125,
+                "Qwen/tri-axis-test",
+                4,
+                8.0,
+                &[a],
+                &[b],
+                &adam_m,
+                &adam_v,
+                &[dynamic],
+                &[LoraTpShardLayout::RoutedExpertFusedGateUp],
+                std::slice::from_ref(&identity),
+                &topology,
+            )
+            .unwrap();
+        }
+    }
+
+    fn rank_for_coordinates(tensor: usize, expert: usize, data: usize) -> usize {
+        (0..8)
+            .find(|global_rank| {
+                let coordinates = tp_ep_dp_topology(*global_rank).coordinates;
+                coordinates.tensor == tensor
+                    && coordinates.expert == expert
+                    && coordinates.data == data
+            })
+            .expect("tri-axis coordinates must map to one global rank")
+    }
+
+    fn replace_checkpoint_tensor(path: &Path, name: &str, value: f64) {
+        let mut tensors = read_named_tensors(path).unwrap();
+        let original = tensors.get(name).unwrap();
+        tensors.insert(
+            name.to_string(),
+            Tensor::full(original.size(), value, (original.kind(), original.device())),
+        );
+        save_named_tensors(path, tensors.into_iter().collect()).unwrap();
+    }
+
+    fn write_test_manifest_and_receipt(rank_dir: &Path, manifest: &CheckpointManifest) {
+        let manifest_contents = serde_json::to_vec_pretty(manifest).unwrap();
+        write_atomic(&rank_dir.join("manifest.json"), &manifest_contents).unwrap();
+        let receipt = checkpoint_rank_receipt(manifest, &manifest_contents).unwrap();
+        write_atomic(
+            &rank_dir.join(RANK_RECEIPT_FILE),
+            &serde_json::to_vec_pretty(&receipt).unwrap(),
+        )
+        .unwrap();
+    }
+
+    fn refresh_checkpoint_file_digest(path: &Path) {
+        let rank_dir = path.parent().unwrap();
+        let file = path.file_name().unwrap().to_str().unwrap();
+        let manifest_path = rank_dir.join("manifest.json");
+        let mut manifest = read_checkpoint_manifest(&manifest_path).unwrap();
+        manifest
+            .file_digests
+            .insert(file.to_string(), stable_file_digest(path).unwrap());
+        write_test_manifest_and_receipt(rank_dir, &manifest);
+    }
+
+    #[test]
+    fn tp_ep_dp_v5_full_rank_set_resumes_every_rank() {
+        let dir = tempfile::tempdir().unwrap();
+        save_tp_ep_dp_checkpoint(dir.path());
+
+        for global_rank in 0..8 {
+            let topology = tp_ep_dp_topology(global_rank);
+            let loaded = load_checkpoint_for_topology(dir.path(), &topology).unwrap();
+            let value = (topology.coordinates.expert * 10 + topology.coordinates.tensor) as f64;
+            assert_eq!(loaded.manifest.step, 31);
+            assert_eq!(loaded.lora_a[0].double_value(&[0, 0, 0]), value + 1.0);
+            assert_eq!(loaded.adam_m[1].double_value(&[0, 0, 0]), value + 4.0);
+            assert_eq!(
+                loaded.dynamic_adapters[0].adam_v[1].double_value(&[0, 0, 0]),
+                value + 16.0
+            );
+            assert!(loaded
+                .manifest
+                .tensor_shards
+                .iter()
+                .all(|shard| { shard.replicated_axes.contains(&ParallelAxis::Data) }));
+        }
+    }
+
+    #[test]
+    fn compact_receipts_avoid_remote_manifest_parsing_and_fail_closed() {
+        let dir = tempfile::tempdir().unwrap();
+        save_tp_ep_dp_checkpoint(dir.path());
+        let remote_manifest = rank_checkpoint_dir(dir.path(), 1).join("manifest.json");
+        std::fs::write(&remote_manifest, b"not-json").unwrap();
+
+        load_checkpoint_for_topology(dir.path(), &tp_ep_dp_topology(0))
+            .expect("rank 0 preflight should consume only the remote compact receipt");
+        let error = load_checkpoint_for_topology(dir.path(), &tp_ep_dp_topology(1))
+            .err()
+            .expect("the owner rank must validate its full local manifest");
+        assert!(error.to_string().contains("parse"));
+
+        let missing_receipt = tempfile::tempdir().unwrap();
+        save_tp_ep_dp_checkpoint(missing_receipt.path());
+        let receipt = rank_checkpoint_dir(missing_receipt.path(), 7).join(RANK_RECEIPT_FILE);
+        std::fs::rename(&receipt, receipt.with_extension("missing")).unwrap();
+        let error = load_checkpoint_for_topology(missing_receipt.path(), &tp_ep_dp_topology(0))
+            .err()
+            .expect("new v5 checkpoints must not fall back when a receipt is missing");
+        assert!(error.to_string().contains(RANK_RECEIPT_FILE));
+    }
+
+    #[test]
+    fn digest_v5_without_receipts_retains_full_manifest_compatibility() {
+        let dir = tempfile::tempdir().unwrap();
+        save_tp_ep_dp_checkpoint(dir.path());
+        for global_rank in 0..8 {
+            let manifest_path = rank_checkpoint_dir(dir.path(), global_rank).join("manifest.json");
+            let mut manifest = read_checkpoint_manifest(&manifest_path).unwrap();
+            manifest.rank_receipt_version = None;
+            write_atomic(
+                &manifest_path,
+                &serde_json::to_vec_pretty(&manifest).unwrap(),
+            )
+            .unwrap();
+        }
+
+        load_checkpoint_for_topology(dir.path(), &tp_ep_dp_topology(0))
+            .expect("digest v5 checkpoints written before compact receipts must remain readable");
+    }
+
+    #[test]
+    fn tp_ep_dp_v5_resume_rejects_divergent_adapter_replica() {
+        let dir = tempfile::tempdir().unwrap();
+        save_tp_ep_dp_checkpoint(dir.path());
+        let corrupt_rank = rank_for_coordinates(0, 0, 1);
+        let corrupt_path =
+            rank_checkpoint_dir(dir.path(), corrupt_rank).join("adapter.safetensors");
+        replace_checkpoint_tensor(&corrupt_path, "a_0", 999.0);
+        refresh_checkpoint_file_digest(&corrupt_path);
+
+        let error = load_checkpoint_for_topology(dir.path(), &tp_ep_dp_topology(0))
+            .err()
+            .expect("divergent DP adapter replica must fail preflight");
+        assert!(error
+            .to_string()
+            .contains("data-parallel replica content digest differs"));
+        assert!(error.to_string().contains("adapter.safetensors"));
+    }
+
+    #[test]
+    fn tp_ep_dp_v5_resume_rejects_divergent_optimizer_replica() {
+        let dir = tempfile::tempdir().unwrap();
+        save_tp_ep_dp_checkpoint(dir.path());
+        let corrupt_rank = rank_for_coordinates(1, 1, 1);
+        let corrupt_path =
+            rank_checkpoint_dir(dir.path(), corrupt_rank).join("optimizer.safetensors");
+        replace_checkpoint_tensor(&corrupt_path, "b_1", 777.0);
+        refresh_checkpoint_file_digest(&corrupt_path);
+
+        let error = load_checkpoint_for_topology(dir.path(), &tp_ep_dp_topology(0))
+            .err()
+            .expect("divergent DP optimizer replica must fail preflight");
+        assert!(error
+            .to_string()
+            .contains("data-parallel replica content digest differs"));
+        assert!(error.to_string().contains("optimizer.safetensors"));
+    }
+
+    #[test]
+    fn tp_ep_dp_v5_resume_rejects_tensor_file_newer_than_manifest() {
+        let dir = tempfile::tempdir().unwrap();
+        save_tp_ep_dp_checkpoint(dir.path());
+        let corrupt_rank = rank_for_coordinates(0, 1, 1);
+        let corrupt_path =
+            rank_checkpoint_dir(dir.path(), corrupt_rank).join("adapter.safetensors");
+        replace_checkpoint_tensor(&corrupt_path, "a_0", 555.0);
+
+        let error = load_checkpoint_for_topology(dir.path(), &tp_ep_dp_topology(corrupt_rank))
+            .err()
+            .expect("stale manifest must not accept newly overwritten tensor content");
+        assert!(error
+            .to_string()
+            .contains("checkpoint content digest differs from manifest"));
+    }
+
+    #[test]
+    fn tp_ep_dp_v5_resume_rejects_missing_rank_and_mixed_generation() {
+        let missing = tempfile::tempdir().unwrap();
+        save_tp_ep_dp_checkpoint(missing.path());
+        let missing_manifest = rank_checkpoint_dir(missing.path(), 7).join("manifest.json");
+        std::fs::rename(
+            &missing_manifest,
+            missing_manifest.with_extension("missing"),
+        )
+        .unwrap();
+        let error = load_checkpoint_for_topology(missing.path(), &tp_ep_dp_topology(0))
+            .err()
+            .expect("missing rank manifest must fail preflight");
+        assert!(error.to_string().contains("rank-00007/manifest.json"));
+
+        let mixed = tempfile::tempdir().unwrap();
+        save_tp_ep_dp_checkpoint(mixed.path());
+        let mixed_manifest_path = rank_checkpoint_dir(mixed.path(), 6).join("manifest.json");
+        let mut mixed_manifest = read_checkpoint_manifest(&mixed_manifest_path).unwrap();
+        mixed_manifest.step += 1;
+        write_test_manifest_and_receipt(mixed_manifest_path.parent().unwrap(), &mixed_manifest);
+        let error = load_checkpoint_for_topology(mixed.path(), &tp_ep_dp_topology(0))
+            .err()
+            .expect("mixed checkpoint generation must fail preflight");
+        assert!(error
+            .to_string()
+            .contains("different checkpoint generation"));
+    }
+
+    fn tp_state(value: f64) -> (Vec<Tensor>, Vec<Tensor>, Vec<Tensor>, Vec<Tensor>) {
+        let a = Tensor::full([2, 3], value, (tch::Kind::Float, tch::Device::Cpu));
+        let b = Tensor::full([4, 2], value + 1.0, (tch::Kind::Float, tch::Device::Cpu));
+        let m = vec![a.zeros_like(), b.zeros_like()];
+        let v = vec![a.ones_like(), b.ones_like()];
+        (vec![a], vec![b], m, v)
+    }
+
+    #[test]
+    fn expert_parallel_checkpoint_uses_rank_scoped_directories() {
+        let dir = tempfile::tempdir().unwrap();
+        for rank in 0..2 {
+            let topology = ep_topology(rank, 2);
+            let (a, b, m, v) = tp_state(rank as f64 + 1.0);
+            save_checkpoint_with_dynamic_for_topology(
+                dir.path(),
+                7,
+                0.25,
+                "Qwen/test",
+                2,
+                4.0,
+                &a,
+                &b,
+                &m,
+                &v,
+                &[],
+                &[],
+                &tp_fixed_identities(),
+                &topology,
+            )
+            .unwrap();
+        }
+
+        assert!(dir.path().join("rank-00000/manifest.json").is_file());
+        assert!(dir.path().join("rank-00001/manifest.json").is_file());
+        assert!(!dir.path().join("manifest.json").exists());
+        let rank1 = load_checkpoint_for_topology(dir.path(), &ep_topology(1, 2)).unwrap();
+        assert_eq!(rank1.lora_a[0].double_value(&[0, 0]), 2.0);
+    }
+
+    #[test]
+    fn checkpoint_topology_preserves_custom_rank_order_and_coordinates() {
+        let topology = ParallelTopology::with_order(2, 1, 1, 2, 1, "ep-tp").unwrap();
+        let manifest = ParallelCheckpointManifest::from_topology(4, 1, &topology).unwrap();
+        assert_eq!(manifest.rank_order, topology.order());
+        assert_eq!(manifest.coordinates.tensor, 0);
+        assert_eq!(manifest.coordinates.expert, 1);
+        assert_eq!(manifest.tensor_model_parallel_rank, 0);
+    }
+
+    #[test]
+    fn routed_expert_shards_record_ep_and_tp_placements() {
+        let topology = tp_ep_topology(3);
+        let gate_up_a = Tensor::zeros([2, 4, 8], (tch::Kind::Float, tch::Device::Cpu));
+        let gate_up_b = Tensor::zeros([2, 6, 4], (tch::Kind::Float, tch::Device::Cpu));
+        let down_a = Tensor::zeros([2, 4, 5], (tch::Kind::Float, tch::Device::Cpu));
+        let down_b = Tensor::zeros([2, 8, 4], (tch::Kind::Float, tch::Device::Cpu));
+
+        let gate_up_a = tensor_shard(
+            &topology,
+            None,
+            4,
+            "adapter.safetensors",
+            "a_0".to_string(),
+            "lora_a",
+            LoraSide::A,
+            LoraTpShardLayout::RoutedExpertFusedGateUp,
+            &gate_up_a,
+        )
+        .unwrap();
+        let gate_up_b = tensor_shard(
+            &topology,
+            None,
+            4,
+            "adapter.safetensors",
+            "b_0".to_string(),
+            "lora_b",
+            LoraSide::B,
+            LoraTpShardLayout::RoutedExpertFusedGateUp,
+            &gate_up_b,
+        )
+        .unwrap();
+        let down_a = tensor_shard(
+            &topology,
+            None,
+            4,
+            "adapter.safetensors",
+            "a_1".to_string(),
+            "lora_a",
+            LoraSide::A,
+            LoraTpShardLayout::RoutedExpertDown,
+            &down_a,
+        )
+        .unwrap();
+        let down_b = tensor_shard(
+            &topology,
+            None,
+            4,
+            "adapter.safetensors",
+            "b_1".to_string(),
+            "lora_b",
+            LoraSide::B,
+            LoraTpShardLayout::RoutedExpertDown,
+            &down_b,
+        )
+        .unwrap();
+
+        assert_eq!(gate_up_a.global_shape, vec![4, 4, 8]);
+        assert_eq!(gate_up_a.placements.len(), 1);
+        assert_eq!(gate_up_a.placements[0].parallel_axis, ParallelAxis::Expert);
+        assert!(gate_up_a.replicated_axes.contains(&ParallelAxis::Tensor));
+
+        assert_eq!(gate_up_b.global_shape, vec![4, 12, 4]);
+        assert_eq!(gate_up_b.global_offset, vec![2, 0, 0]);
+        assert_eq!(gate_up_b.placements.len(), 2);
+        let tp = gate_up_b
+            .placements
+            .iter()
+            .find(|placement| placement.parallel_axis == ParallelAxis::Tensor)
+            .unwrap();
+        assert_eq!(tp.tensor_axis, 1);
+        assert_eq!(tp.segments.len(), 2);
+        assert_eq!(tp.segments[0].global_offset, 3);
+        assert_eq!(tp.segments[1].global_offset, 9);
+
+        assert_eq!(down_a.global_shape, vec![4, 4, 10]);
+        assert_eq!(down_a.global_offset, vec![2, 0, 5]);
+        assert_eq!(down_a.placements.len(), 2);
+        assert_eq!(down_b.global_shape, vec![4, 8, 4]);
+        assert_eq!(down_b.placements.len(), 1);
+        assert!(down_b.replicated_axes.contains(&ParallelAxis::Tensor));
+    }
+
+    #[test]
+    fn routed_expert_projection_layout_does_not_shard_lora_rank() {
+        let topology = tp_ep_topology(0);
+        let tensor = Tensor::zeros([2, 3, 8], (tch::Kind::Float, tch::Device::Cpu));
+        let shard = tensor_shard(
+            &topology,
+            None,
+            3,
+            "adapter.safetensors",
+            "a_0".to_string(),
+            "lora_a",
+            LoraSide::A,
+            LoraTpShardLayout::RoutedExpertFusedGateUp,
+            &tensor,
+        )
+        .unwrap();
+        assert_eq!(shard.global_lora_rank, 3);
+        assert_eq!(shard.global_shape, vec![4, 3, 8]);
+    }
+
+    #[test]
+    fn distributed_merge_reconstructs_fused_expert_gate_up_order() {
+        let dir = tempfile::tempdir().unwrap();
+        let identity = LoraSlotIdentity {
+            index: 0,
+            layer: 0,
+            module: "experts_gate_up_proj".to_string(),
+        };
+        for global_rank in 0..4 {
+            let topology = tp_ep_topology(global_rank);
+            let tp_rank = topology.coordinates.tensor as i64;
+            let ep_rank = topology.coordinates.expert as i64;
+            let a_values = (0..2)
+                .flat_map(|local_expert| {
+                    let expert = ep_rank * 2 + local_expert;
+                    (0..4).flat_map(move |rank| {
+                        (0..2).map(move |hidden| (expert * 100 + rank * 10 + hidden) as f32)
+                    })
+                })
+                .collect::<Vec<_>>();
+            let local_rows = [
+                tp_rank * 2,
+                tp_rank * 2 + 1,
+                4 + tp_rank * 2,
+                5 + tp_rank * 2,
+            ];
+            let b_values = (0..2)
+                .flat_map(|local_expert| {
+                    let expert = ep_rank * 2 + local_expert;
+                    local_rows.into_iter().flat_map(move |row| {
+                        (0..4).map(move |rank| (expert * 1000 + row * 10 + rank) as f32)
+                    })
+                })
+                .collect::<Vec<_>>();
+            let a = Tensor::from_slice(&a_values).reshape([2, 4, 2]);
+            let b = Tensor::from_slice(&b_values).reshape([2, 4, 4]);
+            let adam_m = vec![a.zeros_like(), b.zeros_like()];
+            let adam_v = vec![a.zeros_like(), b.zeros_like()];
+            save_checkpoint_with_dynamic_for_topology(
+                dir.path(),
+                11,
+                0.5,
+                "Qwen/test",
+                4,
+                8.0,
+                &[a],
+                &[b],
+                &adam_m,
+                &adam_v,
+                &[],
+                &[LoraTpShardLayout::RoutedExpertFusedGateUp],
+                std::slice::from_ref(&identity),
+                &topology,
+            )
+            .unwrap();
+        }
+
+        for global_rank in 0..4 {
+            let topology = tp_ep_topology(global_rank);
+            let restored = load_checkpoint_for_topology(dir.path(), &topology).unwrap();
+            assert_eq!(restored.manifest.step, 11);
+            assert_eq!(restored.lora_a[0].size(), [2, 4, 2]);
+            assert_eq!(restored.lora_b[0].size(), [2, 4, 4]);
+        }
+
+        let merged = merge_distributed_adapter_checkpoint(dir.path(), None).unwrap();
+        assert_eq!(merged.step, 11);
+        assert_eq!(merged.slots.len(), 1);
+        assert_eq!(merged.slots[0].lora_a.size(), [4, 4, 2]);
+        assert_eq!(merged.slots[0].lora_b.size(), [4, 8, 4]);
+        for expert in 0..4 {
+            for row in 0..8 {
+                for rank in 0..4 {
+                    assert_eq!(
+                        merged.slots[0].lora_b.double_value(&[expert, row, rank]),
+                        (expert * 1000 + row * 10 + rank) as f64
+                    );
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn distributed_merge_preserves_dynamic_tenant_identity_and_clock() {
+        let dir = tempfile::tempdir().unwrap();
+        for global_rank in 0..2 {
+            let topology = tp_topology(global_rank, 2);
+            let (a, b, m, v) = tp_state(global_rank as f64 + 1.0);
+            let dynamic = tp_dynamic_adapter(global_rank as f64 + 10.0);
+            save_checkpoint_with_dynamic_for_topology(
+                dir.path(),
+                7,
+                0.25,
+                "Qwen/test",
+                4,
+                8.0,
+                &a,
+                &b,
+                &m,
+                &v,
+                &[dynamic],
+                &[],
+                &tp_fixed_identities(),
+                &topology,
+            )
+            .unwrap();
+        }
+
+        let merged = merge_distributed_adapter_checkpoint(dir.path(), Some(9)).unwrap();
+        assert_eq!(merged.rank, 6);
+        assert_eq!(merged.optimizer_step, 4);
+        assert_eq!(merged.slots.len(), 1);
+        assert_eq!(merged.slots[0].identity.module, "q_proj");
+        assert_eq!(merged.slots[0].lora_a.size(), [6, 5]);
+        assert_eq!(merged.slots[0].lora_b.size(), [7, 6]);
+        assert_eq!(merged.slots[0].lora_a.double_value(&[0, 0]), 10.0);
+        assert_eq!(merged.slots[0].lora_a.double_value(&[3, 0]), 11.0);
+    }
+
+    fn tp_fixed_identities() -> [LoraSlotIdentity; 1] {
+        [LoraSlotIdentity {
+            index: 0,
+            layer: 0,
+            module: "in_proj_qkv".to_string(),
+        }]
+    }
+
+    fn tp_dynamic_adapter(value: f64) -> DynamicAdapterCheckpoint {
+        let a = Tensor::full([3, 5], value, (tch::Kind::Float, tch::Device::Cpu));
+        let b = Tensor::full([7, 3], value + 1.0, (tch::Kind::Float, tch::Device::Cpu));
+        DynamicAdapterCheckpoint {
+            manifest: DynamicAdapterManifest {
+                id: 9,
+                rank: 6,
+                alpha: 12.0,
+                optimizer_step: 4,
+                optimizer_lr: Some(5e-4),
+                optimizer_beta1: None,
+                optimizer_beta2: None,
+                optimizer_eps: None,
+                target_layers: vec![1],
+                target_modules: vec!["q_proj".into()],
+                shard_layouts: Vec::new(),
+                slot_identities: vec![LoraSlotIdentity {
+                    index: 0,
+                    layer: 1,
+                    module: "q_proj".to_string(),
+                }],
+                parameter_count: 1,
+                optimizer_count: 2,
+            },
+            lora_a: vec![a.shallow_clone()],
+            lora_b: vec![b.shallow_clone()],
+            adam_m: vec![a.zeros_like(), b.zeros_like()],
+            adam_v: vec![a.ones_like(), b.ones_like()],
+        }
+    }
+
+    #[test]
+    fn tensor_parallel_ranks_use_distinct_paths_and_resume_same_topology() {
+        let dir = tempfile::tempdir().unwrap();
+        for global_rank in 0..2 {
+            let topology = tp_topology(global_rank, 2);
+            let (a, b, m, v) = tp_state(global_rank as f64 + 1.0);
+            let dynamic = tp_dynamic_adapter(global_rank as f64 + 10.0);
+            save_checkpoint_with_dynamic_for_topology(
+                dir.path(),
+                7,
+                0.25,
+                "Qwen/test",
+                4,
+                8.0,
+                &a,
+                &b,
+                &m,
+                &v,
+                &[dynamic],
+                &[],
+                &tp_fixed_identities(),
+                &topology,
+            )
+            .unwrap();
+        }
+
+        let rank0_dir = dir.path().join("rank-00000");
+        let rank1_dir = dir.path().join("rank-00001");
+        assert!(rank0_dir.join("manifest.json").is_file());
+        assert!(rank1_dir.join("manifest.json").is_file());
+
+        let topology = tp_topology(1, 2);
+        let loaded = load_checkpoint_for_topology(dir.path(), &topology).unwrap();
+        assert_eq!(loaded.manifest.format, TP_CHECKPOINT_FORMAT);
+        assert_eq!(loaded.manifest.parallel.as_ref(), Some(&topology));
+        assert_eq!(loaded.manifest.tensor_shards.len(), 12);
+        assert_eq!(loaded.lora_a[0].double_value(&[0, 0]), 2.0);
+        assert_eq!(loaded.dynamic_adapters.len(), 1);
+        assert_eq!(
+            loaded.dynamic_adapters[0].lora_a[0].double_value(&[0, 0]),
+            11.0
+        );
+
+        let a_shard = loaded
+            .manifest
+            .tensor_shards
+            .iter()
+            .find(|shard| shard.state == "lora_a")
+            .unwrap();
+        assert_eq!(a_shard.global_lora_rank, 4);
+        assert_eq!(a_shard.local_shape, vec![2, 3]);
+        assert_eq!(a_shard.global_shape, vec![4, 3]);
+        assert_eq!(a_shard.partition_axis, 0);
+        assert_eq!(a_shard.global_offset, vec![2, 0]);
+        assert_eq!(a_shard.replica_identity, "global-rank-1-tp-rank-1");
+
+        let b_shard = loaded
+            .manifest
+            .tensor_shards
+            .iter()
+            .find(|shard| shard.state == "lora_b")
+            .unwrap();
+        assert_eq!(b_shard.global_shape, vec![4, 4]);
+        assert_eq!(b_shard.partition_axis, 1);
+        assert_eq!(b_shard.global_offset, vec![0, 2]);
+
+        let dynamic_a_shard = loaded
+            .manifest
+            .tensor_shards
+            .iter()
+            .find(|shard| shard.state == "lora_a" && shard.adapter_id == Some(9))
+            .unwrap();
+        assert_eq!(dynamic_a_shard.global_lora_rank, 6);
+        assert_eq!(dynamic_a_shard.global_shape, vec![6, 5]);
+        assert_eq!(dynamic_a_shard.global_offset, vec![3, 0]);
+    }
+
+    #[test]
+    fn tensor_parallel_resume_rejects_different_tp_size() {
+        let dir = tempfile::tempdir().unwrap();
+        let topology = tp_topology(0, 2);
+        let (a, b, m, v) = tp_state(1.0);
+        save_checkpoint_with_dynamic_for_topology(
+            dir.path(),
+            7,
+            0.25,
+            "Qwen/test",
+            4,
+            8.0,
+            &a,
+            &b,
+            &m,
+            &v,
+            &[],
+            &[],
+            &tp_fixed_identities(),
+            &topology,
+        )
+        .unwrap();
+
+        let different_tp = tp_topology(0, 4);
+        let error = load_checkpoint_for_topology(dir.path(), &different_tp)
+            .err()
+            .expect("different TP size must fail");
+        assert!(error.to_string().contains("topology mismatch"));
+    }
+
+    #[test]
+    fn tensor_parallel_resume_rejects_another_ranks_shard() {
+        let dir = tempfile::tempdir().unwrap();
+        let rank0 = tp_topology(0, 2);
+        let (a, b, m, v) = tp_state(1.0);
+        save_checkpoint_with_dynamic_for_topology(
+            dir.path(),
+            7,
+            0.25,
+            "Qwen/test",
+            4,
+            8.0,
+            &a,
+            &b,
+            &m,
+            &v,
+            &[],
+            &[],
+            &tp_fixed_identities(),
+            &rank0,
+        )
+        .unwrap();
+
+        let rank1_dir = dir.path().join("rank-00001");
+        std::fs::create_dir(&rank1_dir).unwrap();
+        for file in [
+            "manifest.json",
+            RANK_RECEIPT_FILE,
+            "adapter.safetensors",
+            "optimizer.safetensors",
+        ] {
+            std::fs::copy(
+                dir.path().join("rank-00000").join(file),
+                rank1_dir.join(file),
+            )
+            .unwrap();
+        }
+
+        let rank1 = tp_topology(1, 2);
+        let error = load_checkpoint_for_topology(dir.path(), &rank1)
+            .err()
+            .expect("loading another rank's shard must fail");
+        assert!(error.to_string().contains("topology mismatch"));
+    }
+
+    #[test]
+    fn tensor_parallel_projection_layouts_record_global_tensor_geometry() {
+        let dir = tempfile::tempdir().unwrap();
+        let column_a = Tensor::zeros([4, 8], (tch::Kind::Float, tch::Device::Cpu));
+        let column_b = Tensor::zeros([8, 4], (tch::Kind::Float, tch::Device::Cpu));
+        let flat_qkv_a = Tensor::zeros([4, 8], (tch::Kind::Float, tch::Device::Cpu));
+        let flat_qkv_b = Tensor::zeros([12, 4], (tch::Kind::Float, tch::Device::Cpu));
+        let row_a = Tensor::zeros([4, 4], (tch::Kind::Float, tch::Device::Cpu));
+        let row_b = Tensor::zeros([8, 4], (tch::Kind::Float, tch::Device::Cpu));
+        let lora_a = vec![
+            column_a.shallow_clone(),
+            flat_qkv_a.shallow_clone(),
+            row_a.shallow_clone(),
+        ];
+        let lora_b = vec![
+            column_b.shallow_clone(),
+            flat_qkv_b.shallow_clone(),
+            row_b.shallow_clone(),
+        ];
+        let adam_m = vec![
+            column_a.zeros_like(),
+            column_b.zeros_like(),
+            flat_qkv_a.zeros_like(),
+            flat_qkv_b.zeros_like(),
+            row_a.zeros_like(),
+            row_b.zeros_like(),
+        ];
+        let adam_v = vec![
+            column_a.ones_like(),
+            column_b.ones_like(),
+            flat_qkv_a.ones_like(),
+            flat_qkv_b.ones_like(),
+            row_a.ones_like(),
+            row_b.ones_like(),
+        ];
+        let layouts = [
+            LoraTpShardLayout::ColumnParallel,
+            LoraTpShardLayout::FlatQkvColumnParallel {
+                q_rows: 4,
+                k_rows: 4,
+                v_rows: 16,
+            },
+            LoraTpShardLayout::RowParallel,
+        ];
+        let identities = [
+            LoraSlotIdentity {
+                index: 0,
+                layer: 0,
+                module: "q_proj".to_string(),
+            },
+            LoraSlotIdentity {
+                index: 4,
+                layer: 1,
+                module: "in_proj_qkv".to_string(),
+            },
+            LoraSlotIdentity {
+                index: 3,
+                layer: 0,
+                module: "o_proj".to_string(),
+            },
+        ];
+
+        for global_rank in 0..2 {
+            save_checkpoint_with_dynamic_for_topology(
+                dir.path(),
+                3,
+                0.5,
+                "Qwen/test",
+                4,
+                8.0,
+                &lora_a,
+                &lora_b,
+                &adam_m,
+                &adam_v,
+                &[],
+                &layouts,
+                &identities,
+                &tp_topology(global_rank, 2),
+            )
+            .unwrap();
+        }
+
+        let topology = tp_topology(1, 2);
+        let loaded = load_checkpoint_for_topology(dir.path(), &topology).unwrap();
+        assert_eq!(loaded.manifest.fixed_shard_layouts, layouts);
+        assert_eq!(loaded.manifest.fixed_slot_identities, identities);
+        let shard = |name: &str| {
+            loaded
+                .manifest
+                .tensor_shards
+                .iter()
+                .find(|shard| shard.file == "adapter.safetensors" && shard.tensor_name == name)
+                .unwrap()
+        };
+
+        let column_a_shard = shard("a_0");
+        assert_eq!(column_a_shard.layout, LoraTpShardLayout::ColumnParallel);
+        assert!(column_a_shard.replicated);
+        assert_eq!(column_a_shard.global_shape, vec![4, 8]);
+        assert_eq!(column_a_shard.global_offset, vec![0, 0]);
+        assert_eq!(column_a_shard.replica_identity, "tp-replicated");
+
+        let column_b_shard = shard("b_0");
+        assert!(!column_b_shard.replicated);
+        assert_eq!(column_b_shard.partition_axis, 0);
+        assert_eq!(column_b_shard.global_shape, vec![16, 4]);
+        assert_eq!(column_b_shard.global_offset, vec![8, 0]);
+
+        let flat_qkv_a_shard = shard("a_1");
+        assert_eq!(
+            flat_qkv_a_shard.layout,
+            LoraTpShardLayout::FlatQkvColumnParallel {
+                q_rows: 4,
+                k_rows: 4,
+                v_rows: 16,
+            }
+        );
+        assert!(flat_qkv_a_shard.replicated);
+        assert_eq!(flat_qkv_a_shard.global_shape, vec![4, 8]);
+        assert_eq!(flat_qkv_a_shard.global_offset, vec![0, 0]);
+
+        let flat_qkv_b_shard = shard("b_1");
+        assert_eq!(
+            flat_qkv_b_shard.layout,
+            LoraTpShardLayout::FlatQkvColumnParallel {
+                q_rows: 4,
+                k_rows: 4,
+                v_rows: 16,
+            }
+        );
+        assert!(!flat_qkv_b_shard.replicated);
+        assert_eq!(flat_qkv_b_shard.partition_axis, 0);
+        assert_eq!(flat_qkv_b_shard.global_shape, vec![24, 4]);
+        assert_eq!(flat_qkv_b_shard.global_offset, vec![0, 0]);
+        assert_eq!(
+            flat_qkv_b_shard.segments,
+            vec![
+                TensorShardSegmentManifest {
+                    local_offset: 0,
+                    global_offset: 2,
+                    length: 2,
+                },
+                TensorShardSegmentManifest {
+                    local_offset: 2,
+                    global_offset: 6,
+                    length: 2,
+                },
+                TensorShardSegmentManifest {
+                    local_offset: 4,
+                    global_offset: 16,
+                    length: 8,
+                },
+            ]
+        );
+
+        let row_a_shard = shard("a_2");
+        assert_eq!(row_a_shard.layout, LoraTpShardLayout::RowParallel);
+        assert!(!row_a_shard.replicated);
+        assert_eq!(row_a_shard.partition_axis, 1);
+        assert_eq!(row_a_shard.global_shape, vec![4, 8]);
+        assert_eq!(row_a_shard.global_offset, vec![0, 4]);
+
+        let row_b_shard = shard("b_2");
+        assert!(row_b_shard.replicated);
+        assert_eq!(row_b_shard.global_shape, vec![8, 4]);
+        assert_eq!(row_b_shard.global_offset, vec![0, 0]);
+        assert_eq!(row_b_shard.replica_identity, "tp-replicated");
+    }
+
+    #[test]
+    fn flat_qkv_segments_reconstruct_global_lora_and_optimizer_rows() {
+        let layout = LoraTpShardLayout::FlatQkvColumnParallel {
+            q_rows: 4,
+            k_rows: 4,
+            v_rows: 16,
+        };
+        let mut reconstructed = [vec![f64::NAN; 24], vec![f64::NAN; 24], vec![f64::NAN; 24]];
+
+        for rank in 0..2 {
+            let topology = tp_topology(rank, 2);
+            let global_rows = (0..24).map(f64::from).collect::<Vec<_>>();
+            let local_rows = [
+                &global_rows[rank * 2..rank * 2 + 2],
+                &global_rows[4 + rank * 2..4 + rank * 2 + 2],
+                &global_rows[8 + rank * 8..8 + rank * 8 + 8],
+            ]
+            .concat();
+
+            for (state_index, (state, delta)) in
+                [("lora_b", 0.0), ("adam_m", 100.0), ("adam_v", 200.0)]
+                    .into_iter()
+                    .enumerate()
+            {
+                let values = local_rows
+                    .iter()
+                    .map(|value| value + delta)
+                    .collect::<Vec<_>>();
+                let tensor = Tensor::from_slice(&values).reshape([12, 1]).repeat([1, 4]);
+                let shard = tensor_shard(
+                    &topology,
+                    None,
+                    4,
+                    "state.safetensors",
+                    state.to_string(),
+                    state,
+                    LoraSide::B,
+                    layout,
+                    &tensor,
+                )
+                .unwrap();
+                for segment in shard.segments {
+                    for offset in 0..segment.length {
+                        reconstructed[state_index][(segment.global_offset + offset) as usize] =
+                            tensor.double_value(&[segment.local_offset + offset, 0]);
+                    }
+                }
+            }
+        }
+
+        for (state_index, delta) in [0.0, 100.0, 200.0].into_iter().enumerate() {
+            let expected = (0..24)
+                .map(|row| f64::from(row) + delta)
+                .collect::<Vec<_>>();
+            assert_eq!(reconstructed[state_index], expected);
+        }
+    }
+
+    #[test]
+    fn tensor_parallel_loader_accepts_legacy_latent_rank_v3_manifest() {
+        let dir = tempfile::tempdir().unwrap();
+        let topology = tp_topology(0, 2);
+        let (a, b, m, v) = tp_state(1.0);
+        save_checkpoint_with_dynamic_for_topology(
+            dir.path(),
+            7,
+            0.25,
+            "Qwen/test",
+            4,
+            8.0,
+            &a,
+            &b,
+            &m,
+            &v,
+            &[],
+            &[],
+            &tp_fixed_identities(),
+            &topology,
+        )
+        .unwrap();
+
+        let manifest_path = dir.path().join("rank-00000/manifest.json");
+        let mut manifest: serde_json::Value =
+            serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap();
+        let object = manifest.as_object_mut().unwrap();
+        object.insert(
+            "format".to_string(),
+            serde_json::Value::String(LEGACY_TP_CHECKPOINT_FORMAT.to_string()),
+        );
+        object.remove("fixed_shard_layouts");
+        object.remove("fixed_slot_identities");
+        let parallel = object.get_mut("parallel").unwrap().as_object_mut().unwrap();
+        parallel.remove("rank_order");
+        parallel.remove("coordinates");
+        for shard in object
+            .get_mut("tensor_shards")
+            .unwrap()
+            .as_array_mut()
+            .unwrap()
+        {
+            let shard = shard.as_object_mut().unwrap();
+            shard.remove("layout");
+            shard.remove("replicated");
+            shard.remove("placements");
+            shard.remove("replicated_axes");
+        }
+        std::fs::write(
+            &manifest_path,
+            serde_json::to_string_pretty(&manifest).unwrap(),
+        )
+        .unwrap();
+
+        let loaded = load_checkpoint_for_topology(dir.path(), &topology).unwrap();
+        assert_eq!(loaded.manifest.format, LEGACY_TP_CHECKPOINT_FORMAT);
+        assert!(loaded.manifest.fixed_shard_layouts.is_empty());
+        assert!(loaded
+            .manifest
+            .tensor_shards
+            .iter()
+            .all(|shard| shard.layout == LoraTpShardLayout::LatentRank && !shard.replicated));
+    }
+
+    #[test]
+    fn tensor_parallel_loader_accepts_projection_aware_v4_without_v5_metadata() {
+        let dir = tempfile::tempdir().unwrap();
+        let topology = tp_topology(0, 2);
+        let (a, b, m, v) = tp_state(1.0);
+        save_checkpoint_with_dynamic_for_topology(
+            dir.path(),
+            7,
+            0.25,
+            "Qwen/test",
+            4,
+            8.0,
+            &a,
+            &b,
+            &m,
+            &v,
+            &[],
+            &[],
+            &tp_fixed_identities(),
+            &topology,
+        )
+        .unwrap();
+
+        let manifest_path = dir.path().join("rank-00000/manifest.json");
+        let mut manifest: serde_json::Value =
+            serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap();
+        let object = manifest.as_object_mut().unwrap();
+        object.insert(
+            "format".to_string(),
+            serde_json::Value::String(PROJECTION_AWARE_TP_CHECKPOINT_FORMAT.to_string()),
+        );
+        let parallel = object.get_mut("parallel").unwrap().as_object_mut().unwrap();
+        parallel.remove("rank_order");
+        parallel.remove("coordinates");
+        for shard in object
+            .get_mut("tensor_shards")
+            .unwrap()
+            .as_array_mut()
+            .unwrap()
+        {
+            let shard = shard.as_object_mut().unwrap();
+            shard.remove("placements");
+            shard.remove("replicated_axes");
+        }
+        std::fs::write(
+            &manifest_path,
+            serde_json::to_string_pretty(&manifest).unwrap(),
+        )
+        .unwrap();
+
+        let loaded = load_checkpoint_for_topology(dir.path(), &topology).unwrap();
+        assert_eq!(
+            loaded.manifest.format,
+            PROJECTION_AWARE_TP_CHECKPOINT_FORMAT
+        );
+        assert!(loaded
+            .manifest
+            .tensor_shards
+            .iter()
+            .all(|shard| shard.placements.is_empty() && shard.replicated_axes.is_empty()));
+    }
+
+    fn resume_validation_manifest(format: &str) -> CheckpointManifest {
+        CheckpointManifest {
+            format: format.to_string(),
+            rank_receipt_version: None,
+            checkpoint_generation: None,
+            step: 0,
+            fixed_optimizer_step: None,
+            loss: 0.0,
+            model_path: "Qwen/test".to_string(),
+            lora_rank: 4,
+            lora_alpha: 8.0,
+            files: Vec::new(),
+            file_digests: BTreeMap::new(),
+            dynamic_adapters: Vec::new(),
+            parallel: None,
+            tensor_shards: Vec::new(),
+            fixed_shard_layouts: vec![LoraTpShardLayout::ColumnParallel],
+            fixed_slot_identities: vec![LoraSlotIdentity {
+                index: 0,
+                layer: 0,
+                module: "q_proj".to_string(),
+            }],
+            fixed_target_layers: Vec::new(),
+            fixed_target_modules: Vec::new(),
+            pipeline_stage: None,
+        }
+    }
+
+    #[test]
+    fn tensor_parallel_resume_rejects_fixed_slot_identity_mismatch() {
+        let manifest = resume_validation_manifest(TP_CHECKPOINT_FORMAT);
+        let expected = [LoraSlotIdentity {
+            index: 1,
+            layer: 0,
+            module: "k_proj".to_string(),
+        }];
+        let error =
+            validate_fixed_tp_resume(&manifest, &[LoraTpShardLayout::ColumnParallel], &expected)
+                .unwrap_err();
+        assert!(error.to_string().contains("slot identities"));
+    }
+
+    #[test]
+    fn legacy_tensor_parallel_resume_rejects_projection_aware_layouts() {
+        let manifest = resume_validation_manifest(LEGACY_TP_CHECKPOINT_FORMAT);
+        let fixed_error =
+            validate_fixed_tp_resume(&manifest, &[LoraTpShardLayout::ColumnParallel], &[])
+                .unwrap_err();
+        assert!(fixed_error
+            .to_string()
+            .contains("fixed projection-aware attention LoRA"));
+        let dynamic_error =
+            validate_dynamic_tp_resume(&manifest, 17, &[], &[LoraTpShardLayout::RowParallel], &[])
+                .unwrap_err();
+        assert!(dynamic_error.to_string().contains("adapter 17"));
+        validate_fixed_tp_resume(&manifest, &[LoraTpShardLayout::LatentRank], &[]).unwrap();
+    }
+
+    #[test]
+    fn fixed_restore_mapping_supports_compact_and_legacy_positional_slots() {
+        assert_eq!(
+            fixed_restore_slot_indices(2, 2, &[1, 4], 7).unwrap(),
+            vec![1, 4]
+        );
+        assert_eq!(
+            fixed_restore_slot_indices(7, 7, &[1, 4], 7).unwrap(),
+            (0..7).collect::<Vec<_>>()
+        );
+        assert!(fixed_restore_slot_indices(2, 1, &[1, 4], 7).is_err());
+        assert!(fixed_restore_slot_indices(3, 3, &[1, 4], 7).is_err());
+    }
+}
diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs
index 0db5edb9..94a0fd94 100644
--- a/crates/rustrain-qwen3-6/src/kernel.rs
+++ b/crates/rustrain-qwen3-6/src/kernel.rs
@@ -3,48 +3,313 @@
 //! TrainingContext, train_step, and adapter export all happen in C++.
 //! Rust only handles: weight loading, data loading, training loop orchestration.
 
+use crate::lora::Qwen36LoraTargetModule;
+use crate::pipeline::PipelineStageLayout;
+use anyhow::{bail, Context, Result};
 use std::ffi::c_void;
 use std::sync::OnceLock;
 use tch::{Kind, Tensor};
-use anyhow::{Result, bail};
 
 // ── dlopen ──
 
 type FnCreateCtx = unsafe extern "C" fn(
-    *mut *mut c_void, i64, *mut c_void, *mut c_void, *mut c_void,
-    *mut c_void, i64, i32, f64, f64, f64, f64, f64, i64, f64,
-    i64, *const i64, i64,
+    *mut *mut c_void,
+    i64,
+    *mut c_void,
+    *mut c_void,
+    *mut c_void,
+    *mut c_void,
+    i64,
+    i64,
+    i64,
+    i32,
+    i32,
+    f64,
+    f64,
+    f64,
+    f64,
+    f64,
+    i64,
+    f64,
+    i64,
+    *const i64,
+    i64,
+    *const i8,
+    i32,
 ) -> *mut c_void;
+type FnKernelAbiVersion = unsafe extern "C" fn() -> i64;
 type FnTrainStep = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void) -> f64;
-type FnTrainMultiLora = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void, i32, i32) -> f64;
+type FnTrainMicroStep =
+    unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void, f64, i32) -> f64;
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct PipelineWindowV1 {
+    pub struct_size: u32,
+    pub version: u32,
+    pub window_id: i64,
+    pub num_microbatches: i64,
+    pub schedule: i32,
+    pub num_chunks: i32,
+    pub flags: i32,
+}
+pub const PIPELINE_WINDOW_FLAG_DYNAMIC_LORA: i32 = 1;
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct PipelineTickV1 {
+    pub struct_size: u32,
+    pub version: u32,
+    pub window_id: i64,
+    pub forward_mb: i64,
+    pub backward_mb: i64,
+    pub chunk_id: i32,
+    pub phase: i32,
+    pub input_ids: *mut c_void,
+    pub target_mask: *mut c_void,
+    pub attention_mask: *mut c_void,
+    pub gradient_scale: f64,
+}
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct PipelineResultV1 {
+    pub struct_size: u32,
+    pub version: u32,
+    pub status: i32,
+    pub completed_fwd: i64,
+    pub completed_bwd: i64,
+    pub in_flight: i64,
+    pub optimizer_step: i64,
+    pub loss: f64,
+}
+type FnPipelineBegin = unsafe extern "C" fn(*mut c_void, *const PipelineWindowV1) -> i32;
+type FnPipelineBeginSelected = unsafe extern "C" fn(
+    *mut c_void,
+    *const PipelineWindowV1,
+    *const i64,
+    i32,
+) -> i32;
+type FnPipelineTick =
+    unsafe extern "C" fn(*mut c_void, *const PipelineTickV1, *mut PipelineResultV1) -> i32;
+type FnPipelineFinish = unsafe extern "C" fn(*mut c_void, i32, *mut PipelineResultV1) -> i32;
+type FnPipelineFinishReport = unsafe extern "C" fn(
+    *mut c_void,
+    i32,
+    *mut PipelineResultV1,
+    *mut f64,
+    i32,
+) -> i32;
+type FnPipelineAbort = unsafe extern "C" fn(*mut c_void) -> i32;
+type FnTrainMultiLoraSelectedV2 = unsafe extern "C" fn(
+    *mut c_void,
+    *mut c_void,
+    *mut c_void,
+    *mut c_void,
+    *const i64,
+    i32,
+) -> f64;
+type FnTrainMultiLoraSelectedV3 = unsafe extern "C" fn(
+    *mut c_void,
+    *mut c_void,
+    *mut c_void,
+    *mut c_void,
+    *const i64,
+    i32,
+    *mut f64,
+    *mut f64,
+    i32,
+) -> i32;
 type FnEvalStep = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void) -> f64;
+type FnEvalMultiLoraSelected = unsafe extern "C" fn(
+    *mut c_void,
+    *mut c_void,
+    *mut c_void,
+    *mut c_void,
+    *const i64,
+    i32,
+    *mut f64,
+    i32,
+) -> i32;
+type FnHostBatchStep =
+    unsafe extern "C" fn(*mut c_void, *const i64, *const i64, *const i64, i64, i64) -> f64;
+type FnEvalMultiLoraHost = unsafe extern "C" fn(
+    *mut c_void,
+    *const i64,
+    *const i64,
+    *const i64,
+    i64,
+    i64,
+    *const i64,
+    i32,
+    *mut f64,
+    i32,
+) -> i32;
+type FnHostMultiLoraStep = unsafe extern "C" fn(
+    *mut c_void,
+    *const i64,
+    *const i64,
+    *const i64,
+    i64,
+    i64,
+    i32,
+    i32,
+    *const i64,
+    i32,
+) -> f64;
+type FnHostMultiLoraReport = unsafe extern "C" fn(
+    *mut c_void,
+    *const i64,
+    *const i64,
+    *const i64,
+    i64,
+    i64,
+    i32,
+    i32,
+    *const i64,
+    i32,
+    *mut f64,
+    *mut f64,
+    i32,
+) -> i32;
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct MultiLoraLossReport {
+    pub aggregate_loss: f64,
+    pub adapter_losses: Vec<f64>,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct DynamicAdamConfig {
+    pub lr: f64,
+    pub beta1: f64,
+    pub beta2: f64,
+    pub eps: f64,
+}
+
+impl DynamicAdamConfig {
+    fn validate(self) -> Result<Self> {
+        if !self.lr.is_finite() || self.lr < 0.0 {
+            bail!("dynamic LoRA optimizer learning rate must be finite and non-negative");
+        }
+        if !self.beta1.is_finite() || !(0.0..1.0).contains(&self.beta1) {
+            bail!("dynamic LoRA optimizer beta1 must be finite and in [0, 1)");
+        }
+        if !self.beta2.is_finite() || !(0.0..1.0).contains(&self.beta2) {
+            bail!("dynamic LoRA optimizer beta2 must be finite and in [0, 1)");
+        }
+        if !((self.beta1 as f32).is_finite() && (self.beta1 as f32) < 1.0)
+            || !((self.beta2 as f32).is_finite() && (self.beta2 as f32) < 1.0)
+        {
+            bail!("dynamic LoRA optimizer betas must be representable as finite FP32 values below 1");
+        }
+        if !self.eps.is_finite() || self.eps < 0.0 {
+            bail!("dynamic LoRA optimizer epsilon must be finite and non-negative");
+        }
+        Ok(self)
+    }
+}
 type FnGetLoraCount = unsafe extern "C" fn(*mut c_void) -> i64;
 type FnGetLoraA = unsafe extern "C" fn(*mut c_void, i64) -> *mut c_void;
 type FnGetLoraB = unsafe extern "C" fn(*mut c_void, i64) -> *mut c_void;
+type FnSetLoraTensor = unsafe extern "C" fn(*mut c_void, i64, i32, *mut c_void) -> i32;
+type FnGetLoraGradAccumulator = unsafe extern "C" fn(*mut c_void, i64, i32) -> *mut c_void;
+type FnAbortGradientAccumulation = unsafe extern "C" fn(*mut c_void) -> i32;
 type FnGetStepCount = unsafe extern "C" fn(*mut c_void) -> i64;
-type FnExportOptimizer = unsafe extern "C" fn(*mut c_void, *mut *mut c_void, *mut *mut c_void, i64) -> i64;
-type FnImportOptimizer = unsafe extern "C" fn(*mut c_void, *mut *mut c_void, *mut *mut c_void, i64) -> i64;
+type FnSetStepCount = unsafe extern "C" fn(*mut c_void, i64) -> i32;
+type FnExportOptimizer =
+    unsafe extern "C" fn(*mut c_void, *mut *mut c_void, *mut *mut c_void, i64) -> i64;
+type FnImportOptimizer =
+    unsafe extern "C" fn(*mut c_void, *mut *mut c_void, *mut *mut c_void, i64) -> i64;
 type FnFreeCtx = unsafe extern "C" fn(*mut c_void);
 type FnGemm = unsafe extern "C" fn(*mut c_void, *mut c_void, i32) -> *mut c_void;
 type FnFreeTensor = unsafe extern "C" fn(*mut c_void);
 type FnSetMtpWeights = unsafe extern "C" fn(
-    *mut c_void,       // ctx_ptr
-    *mut c_void,       // mtp_fc_ptr
-    *mut c_void,       // mtp_pre_fc_norm_emb_ptr
-    *mut c_void,       // mtp_pre_fc_norm_hidden_ptr
-    *mut c_void,       // mtp_norm_ptr
-    *mut *mut c_void,  // mtp_layer_weight_ptrs
-    i64,               // num_mtp_layer_weights
-    *mut c_void,       // mtp_layer_configs_ptr
-    i64,               // num_mtp_layers
-);
+    *mut c_void,      // ctx_ptr
+    *mut c_void,      // mtp_fc_ptr
+    *mut c_void,      // mtp_pre_fc_norm_emb_ptr
+    *mut c_void,      // mtp_pre_fc_norm_hidden_ptr
+    *mut c_void,      // mtp_norm_ptr
+    *mut *mut c_void, // mtp_layer_weight_ptrs
+    i64,              // num_mtp_layer_weights
+    *mut c_void,      // mtp_layer_configs_ptr
+    i64,              // num_mtp_layers
+) -> i32;
 type FnSetCheckpoint = unsafe extern "C" fn(*mut c_void, i32, i64);
 type FnSetNcclComm = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, i32, i32);
 type FnInitNccl = unsafe extern "C" fn(*mut c_void) -> i32;
+type FnInitParallelNccl = unsafe extern "C" fn(
+    *mut c_void,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+    i32,
+) -> i32;
 type FnSetCudaDevice = unsafe extern "C" fn(i32);
+type FnSetPadTokenId = unsafe extern "C" fn(*mut c_void, i64) -> i32;
+type FnSetBaseTpMlp = unsafe extern "C" fn(*mut c_void, i32) -> i32;
+type FnSetMaxGradNorm = unsafe extern "C" fn(*mut c_void, f64) -> i32;
+type FnSetRouterAuxLossCoef = unsafe extern "C" fn(*mut c_void, f64) -> i32;
 type FnAddLora = unsafe extern "C" fn(*mut c_void, i64, f64, *const i64, i64, *const i8) -> i64;
+type FnAddLoraWithOptimizer =
+    unsafe extern "C" fn(*mut c_void, i64, f64, *const i64, i64, *const i8, f64) -> i64;
+type FnAddLoraWithOptimizerV2 = unsafe extern "C" fn(
+    *mut c_void,
+    i64,
+    f64,
+    *const i64,
+    i64,
+    *const i8,
+    f64,
+    f64,
+    f64,
+    f64,
+) -> i64;
 type FnRemoveLora = unsafe extern "C" fn(*mut c_void, i64) -> i32;
 type FnListLora = unsafe extern "C" fn(*mut c_void, *mut i64, i64) -> i64;
+type FnGetAdapterLoraTensor =
+    unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32) -> *mut c_void;
+type FnSetAdapterLoraTensor =
+    unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32, *mut c_void) -> i32;
+type FnSetAdapterId = unsafe extern "C" fn(*mut c_void, i64, i64) -> i32;
+type FnGetAdapterOptimizerTensor =
+    unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32, i32) -> *mut c_void;
+type FnExportAdapterOptimizerTensorCpuV1 = unsafe extern "C" fn(
+    *mut c_void,
+    i64,
+    i64,
+    *const i8,
+    i32,
+    i32,
+    *mut *mut c_void,
+) -> i32;
+type FnImportAdapterOptimizerStateHostV1 = unsafe extern "C" fn(
+    *mut c_void,
+    i64,
+    *const i64,
+    *const *const i8,
+    *const *mut c_void,
+    i64,
+    i64,
+) -> i32;
+type FnGetAdapterOptimizerResidentCountV1 =
+    unsafe extern "C" fn(*mut c_void, i64) -> i64;
+type FnSetAdapterOptimizerTensor =
+    unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32, i32, *mut c_void) -> i32;
+type FnGetAdapterStepCount = unsafe extern "C" fn(*mut c_void, i64) -> i64;
+type FnSetAdapterStepCount = unsafe extern "C" fn(*mut c_void, i64, i64) -> i32;
+type FnValidateAdapterSteps = unsafe extern "C" fn(*mut c_void, *const i64, *const i64, i32) -> i32;
+type FnGetContextHealth = unsafe extern "C" fn(*mut c_void) -> i32;
 
 #[repr(C)]
 pub struct CppLayerConfig {
@@ -75,12 +340,30 @@ pub struct CppLayerConfig {
 struct KernelHandles {
     create_ctx: FnCreateCtx,
     train_step: FnTrainStep,
-    train_multi_lora: FnTrainMultiLora,
+    train_micro_step: FnTrainMicroStep,
+    pipeline_begin: FnPipelineBegin,
+    pipeline_begin_selected: Option<FnPipelineBeginSelected>,
+    pipeline_tick: FnPipelineTick,
+    pipeline_finish: FnPipelineFinish,
+    pipeline_finish_report: Option<FnPipelineFinishReport>,
+    pipeline_abort: FnPipelineAbort,
+    train_multi_lora_selected_v2: FnTrainMultiLoraSelectedV2,
+    train_multi_lora_selected_v3: FnTrainMultiLoraSelectedV3,
     eval_step: FnEvalStep,
+    eval_multi_lora_selected: Option<FnEvalMultiLoraSelected>,
+    train_step_host_i64: FnHostBatchStep,
+    train_multi_lora_host_i64: FnHostMultiLoraStep,
+    train_multi_lora_host_i64_v2: FnHostMultiLoraReport,
+    eval_step_host_i64: FnHostBatchStep,
+    eval_multi_lora_host_i64: Option<FnEvalMultiLoraHost>,
     get_lora_count: FnGetLoraCount,
     get_lora_a: FnGetLoraA,
     get_lora_b: FnGetLoraB,
+    set_lora_tensor: FnSetLoraTensor,
+    get_lora_grad_accumulator: FnGetLoraGradAccumulator,
+    abort_gradient_accumulation: FnAbortGradientAccumulation,
     get_step_count: FnGetStepCount,
+    set_step_count: FnSetStepCount,
     export_optimizer: FnExportOptimizer,
     import_optimizer: FnImportOptimizer,
     free_ctx: FnFreeCtx,
@@ -90,10 +373,34 @@ struct KernelHandles {
     set_checkpoint: FnSetCheckpoint,
     set_nccl_comm: FnSetNcclComm,
     init_nccl: FnInitNccl,
+    init_parallel_nccl: FnInitParallelNccl,
+    attach_parallel_nccl_no_sync: FnInitParallelNccl,
     set_cuda_device: FnSetCudaDevice,
+    set_pad_token_id: Option<FnSetPadTokenId>,
+    set_base_tp_mlp: FnSetBaseTpMlp,
+    set_max_grad_norm: FnSetMaxGradNorm,
+    set_router_aux_loss_coef: FnSetRouterAuxLossCoef,
     add_lora: FnAddLora,
+    add_lora_v2: FnAddLora,
+    add_lora_with_optimizer: Option<FnAddLoraWithOptimizer>,
+    add_lora_with_optimizer_v2: Option<FnAddLoraWithOptimizerV2>,
+    add_lora_for_restore: FnAddLora,
+    add_lora_for_restore_with_optimizer: FnAddLoraWithOptimizer,
+    add_lora_for_restore_with_optimizer_v2: Option<FnAddLoraWithOptimizerV2>,
     remove_lora: FnRemoveLora,
     list_lora: FnListLora,
+    get_adapter_lora_tensor: FnGetAdapterLoraTensor,
+    set_adapter_lora_tensor: FnSetAdapterLoraTensor,
+    set_adapter_id: FnSetAdapterId,
+    get_adapter_optimizer_tensor: FnGetAdapterOptimizerTensor,
+    export_adapter_optimizer_tensor_cpu: FnExportAdapterOptimizerTensorCpuV1,
+    import_adapter_optimizer_state_host: FnImportAdapterOptimizerStateHostV1,
+    get_adapter_optimizer_resident_count: FnGetAdapterOptimizerResidentCountV1,
+    set_adapter_optimizer_tensor: FnSetAdapterOptimizerTensor,
+    get_adapter_step_count: FnGetAdapterStepCount,
+    set_adapter_step_count: FnSetAdapterStepCount,
+    validate_adapter_steps: Option<FnValidateAdapterSteps>,
+    get_context_health: Option<FnGetContextHealth>,
 }
 
 static KERNELS: OnceLock<Option<KernelHandles>> = OnceLock::new();
@@ -115,27 +422,142 @@ unsafe fn load_kernels() -> Option<KernelHandles> {
     let handle = libc::dlopen(lib_name.as_ptr(), libc::RTLD_LAZY | libc::RTLD_NOLOAD);
     let handle = if handle.is_null() {
         let h = libc::dlopen(lib_name.as_ptr(), libc::RTLD_LAZY);
-        if h.is_null() { return None; }
+        if h.is_null() {
+            return None;
+        }
         h
-    } else { handle };
+    } else {
+        handle
+    };
 
     macro_rules! sym {
         ($name:expr) => {{
             let s = CString::new($name).unwrap();
             let p = libc::dlsym(handle, s.as_ptr());
-            if p.is_null() { return None; }
+            if p.is_null() {
+                return None;
+            }
             std::mem::transmute::<*mut c_void, _>(p)
         }};
     }
+    let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version");
+    if abi_version() != 31 {
+        return None;
+    }
+    let add_lora_with_optimizer = {
+        let name = CString::new("qwen36_add_lora_with_optimizer").unwrap();
+        let symbol = libc::dlsym(handle, name.as_ptr());
+        if symbol.is_null() {
+            None
+        } else {
+            Some(std::mem::transmute::<*mut c_void, FnAddLoraWithOptimizer>(
+                symbol,
+            ))
+        }
+    };
+    let add_lora_with_optimizer_v2 = {
+        let name = CString::new("qwen36_add_lora_with_optimizer_v2").unwrap();
+        let symbol = libc::dlsym(handle, name.as_ptr());
+        if symbol.is_null() {
+            None
+        } else {
+            Some(std::mem::transmute::<*mut c_void, FnAddLoraWithOptimizerV2>(symbol))
+        }
+    };
+    let add_lora_for_restore_with_optimizer_v2 = {
+        let name = CString::new("qwen36_add_lora_for_restore_with_optimizer_v2").unwrap();
+        let symbol = libc::dlsym(handle, name.as_ptr());
+        if symbol.is_null() {
+            None
+        } else {
+            Some(std::mem::transmute::<*mut c_void, FnAddLoraWithOptimizerV2>(symbol))
+        }
+    };
     Some(KernelHandles {
-        create_ctx: sym!("qwen36_create_training_context"),
+        create_ctx: sym!("qwen36_create_training_context_v2"),
         train_step: sym!("qwen36_train_step"),
-        train_multi_lora: sym!("qwen36_train_multi_lora"),
+        train_micro_step: sym!("qwen36_train_micro_step"),
+        pipeline_begin: sym!("qwen36_pipeline_begin_v1"),
+        pipeline_begin_selected: {
+            let name = CString::new("qwen36_pipeline_begin_dynamic_selected_v1").unwrap();
+            let symbol = libc::dlsym(handle, name.as_ptr());
+            if symbol.is_null() {
+                None
+            } else {
+                Some(std::mem::transmute::<*mut c_void, FnPipelineBeginSelected>(symbol))
+            }
+        },
+        pipeline_tick: sym!("qwen36_pipeline_tick_v1"),
+        pipeline_finish: sym!("qwen36_pipeline_finish_v1"),
+        pipeline_finish_report: {
+            let name = CString::new("qwen36_pipeline_finish_dynamic_report_v1").unwrap();
+            let symbol = libc::dlsym(handle, name.as_ptr());
+            if symbol.is_null() {
+                None
+            } else {
+                Some(std::mem::transmute::<*mut c_void, FnPipelineFinishReport>(symbol))
+            }
+        },
+        pipeline_abort: sym!("qwen36_pipeline_abort_v1"),
+        train_multi_lora_selected_v2: sym!("qwen36_train_multi_lora_selected_v2"),
+        train_multi_lora_selected_v3: sym!("qwen36_train_multi_lora_selected_v3"),
         eval_step: sym!("qwen36_eval_step"),
+        eval_multi_lora_selected: {
+            let name = CString::new("qwen36_eval_multi_lora_selected_v1").unwrap();
+            let symbol = libc::dlsym(handle, name.as_ptr());
+            if symbol.is_null() {
+                None
+            } else {
+                Some(std::mem::transmute::<*mut c_void, FnEvalMultiLoraSelected>(
+                    symbol,
+                ))
+            }
+        },
+        train_step_host_i64: sym!("qwen36_train_step_host_i64"),
+        train_multi_lora_host_i64: sym!("qwen36_train_multi_lora_host_i64"),
+        train_multi_lora_host_i64_v2: sym!("qwen36_train_multi_lora_host_i64_v2"),
+        eval_step_host_i64: sym!("qwen36_eval_step_host_i64"),
+        eval_multi_lora_host_i64: {
+            let name = CString::new("qwen36_eval_multi_lora_host_i64_v1").unwrap();
+            let symbol = libc::dlsym(handle, name.as_ptr());
+            if symbol.is_null() {
+                None
+            } else {
+                Some(std::mem::transmute::<*mut c_void, FnEvalMultiLoraHost>(
+                    symbol,
+                ))
+            }
+        },
+        validate_adapter_steps: {
+            let name = CString::new("qwen36_validate_adapter_steps_v1").unwrap();
+            let symbol = libc::dlsym(handle, name.as_ptr());
+            if symbol.is_null() {
+                None
+            } else {
+                Some(std::mem::transmute::<*mut c_void, FnValidateAdapterSteps>(
+                    symbol,
+                ))
+            }
+        },
+        get_context_health: {
+            let name = CString::new("qwen36_get_context_health").unwrap();
+            let symbol = libc::dlsym(handle, name.as_ptr());
+            if symbol.is_null() {
+                None
+            } else {
+                Some(std::mem::transmute::<*mut c_void, FnGetContextHealth>(
+                    symbol,
+                ))
+            }
+        },
         get_lora_count: sym!("qwen36_get_lora_count"),
         get_lora_a: sym!("qwen36_get_lora_a"),
         get_lora_b: sym!("qwen36_get_lora_b"),
+        set_lora_tensor: sym!("qwen36_set_lora_tensor"),
+        get_lora_grad_accumulator: sym!("qwen36_get_lora_grad_accumulator"),
+        abort_gradient_accumulation: sym!("qwen36_abort_gradient_accumulation"),
         get_step_count: sym!("qwen36_get_step_count"),
+        set_step_count: sym!("qwen36_set_step_count"),
         export_optimizer: sym!("qwen36_export_optimizer_state"),
         import_optimizer: sym!("qwen36_import_optimizer_state"),
         free_ctx: sym!("qwen36_free_training_context"),
@@ -145,10 +567,46 @@ unsafe fn load_kernels() -> Option<KernelHandles> {
         set_checkpoint: sym!("qwen36_set_checkpoint"),
         set_nccl_comm: sym!("qwen36_set_nccl_comm"),
         init_nccl: sym!("qwen36_init_nccl"),
+        init_parallel_nccl: sym!("qwen36_init_parallel_nccl_v2"),
+        attach_parallel_nccl_no_sync: sym!("qwen36_attach_parallel_nccl_no_sync_v2"),
         set_cuda_device: sym!("qwen36_set_cuda_device"),
+        set_pad_token_id: {
+            let name = CString::new("qwen36_set_pad_token_id").unwrap();
+            let symbol = libc::dlsym(handle, name.as_ptr());
+            if symbol.is_null() {
+                None
+            } else {
+                Some(std::mem::transmute::<*mut c_void, FnSetPadTokenId>(symbol))
+            }
+        },
+        set_base_tp_mlp: sym!("qwen36_set_base_tp_mlp"),
+        set_max_grad_norm: sym!("qwen36_set_max_grad_norm"),
+        set_router_aux_loss_coef: sym!("qwen36_set_router_aux_loss_coef"),
         add_lora: sym!("qwen36_add_lora"),
+        add_lora_v2: sym!("qwen36_add_lora_v2"),
+        add_lora_with_optimizer,
+        add_lora_with_optimizer_v2,
+        add_lora_for_restore: sym!("qwen36_add_lora_for_restore"),
+        add_lora_for_restore_with_optimizer: sym!("qwen36_add_lora_for_restore_with_optimizer"),
+        add_lora_for_restore_with_optimizer_v2,
         remove_lora: sym!("qwen36_remove_lora"),
         list_lora: sym!("qwen36_list_lora"),
+        get_adapter_lora_tensor: sym!("qwen36_get_adapter_lora_tensor"),
+        set_adapter_lora_tensor: sym!("qwen36_set_adapter_lora_tensor"),
+        set_adapter_id: sym!("qwen36_set_adapter_id"),
+        get_adapter_optimizer_tensor: sym!("qwen36_get_adapter_optimizer_tensor"),
+        export_adapter_optimizer_tensor_cpu: sym!(
+            "qwen36_export_adapter_optimizer_tensor_cpu_v1"
+        ),
+        import_adapter_optimizer_state_host: sym!(
+            "qwen36_import_adapter_optimizer_state_host_v1"
+        ),
+        get_adapter_optimizer_resident_count: sym!(
+            "qwen36_get_adapter_optimizer_resident_count_v1"
+        ),
+        set_adapter_optimizer_tensor: sym!("qwen36_set_adapter_optimizer_tensor"),
+        get_adapter_step_count: sym!("qwen36_get_adapter_step_count"),
+        set_adapter_step_count: sym!("qwen36_set_adapter_step_count"),
     })
 }
 
@@ -169,16 +627,293 @@ fn get_ptr(weights: &std::collections::BTreeMap<String, Tensor>, name: &str) ->
     }
 }
 
+/// Return the rank-local frozen vocabulary shard for TP, or `None` when the
+/// tensor is not an embedding or LM-head weight. This runs during CPU loading.
+pub fn shard_vocab_weight_for_tp(
+    name: &str,
+    tensor: &Tensor,
+    vocab_size: i64,
+    tp_size: usize,
+    tp_rank: usize,
+) -> Result<Option<Tensor>> {
+    if !name.ends_with("embed_tokens.weight") && !name.ends_with("lm_head.weight") {
+        return Ok(None);
+    }
+    if tp_size <= 1 || tp_rank >= tp_size {
+        bail!("invalid vocabulary TP shard: tp_rank={tp_rank}, tp_size={tp_size}");
+    }
+    if vocab_size <= 0 || vocab_size % tp_size as i64 != 0 {
+        bail!("vocab_size={vocab_size} is not divisible by TP_SIZE={tp_size}");
+    }
+    let shape = tensor.size();
+    if shape.len() != 2 || shape[0] != vocab_size {
+        bail!("vocabulary TP weight {name} must have shape [{vocab_size}, hidden], got {shape:?}");
+    }
+    let local_vocab_size = vocab_size / tp_size as i64;
+    Ok(Some(
+        tensor
+            .narrow(0, tp_rank as i64 * local_vocab_size, local_vocab_size)
+            .contiguous(),
+    ))
+}
+
+/// Return the rank-local frozen dense MLP shard for TP, or `None` when the
+/// tensor is not one of gate/up/down. This runs during CPU weight loading.
+pub fn shard_dense_mlp_weight_for_tp(
+    name: &str,
+    tensor: &Tensor,
+    tp_size: usize,
+    tp_rank: usize,
+) -> Result<Option<Tensor>> {
+    let dim = if name.ends_with(".mlp.gate_proj.weight") || name.ends_with(".mlp.up_proj.weight") {
+        0
+    } else if name.ends_with(".mlp.down_proj.weight") {
+        1
+    } else {
+        return Ok(None);
+    };
+    if tp_size <= 1 || tp_rank >= tp_size {
+        bail!("invalid dense MLP TP shard: tp_rank={tp_rank}, tp_size={tp_size}");
+    }
+    let full = *tensor
+        .size()
+        .get(dim as usize)
+        .ok_or_else(|| anyhow::anyhow!("base TP MLP weight {name} has no dimension {dim}"))?;
+    if full <= 0 || full % tp_size as i64 != 0 {
+        bail!(
+            "base TP MLP weight {name} dimension {dim}={full} is not divisible by TP_SIZE={tp_size}"
+        );
+    }
+    let shard = full / tp_size as i64;
+    let start = tp_rank as i64 * shard;
+    Ok(Some(tensor.narrow(dim, start, shard).contiguous()))
+}
+
+/// Return the rank-local frozen MoE MLP shard for expert tensor parallelism.
+/// Routed expert tensors may already be narrowed on their leading EP axis.
+/// The packed gate/up layout is `[gate_all | up_all]`, so each half must be
+/// sliced independently before the local halves are concatenated.
+pub fn shard_moe_mlp_weight_for_tp(
+    name: &str,
+    tensor: &Tensor,
+    tp_size: usize,
+    tp_rank: usize,
+) -> Result<Option<Tensor>> {
+    let is_shared_gate_up = name.ends_with(".mlp.shared_expert.gate_proj.weight")
+        || name.ends_with(".mlp.shared_expert.up_proj.weight");
+    let is_shared_down = name.ends_with(".mlp.shared_expert.down_proj.weight");
+    let is_expert_gate_up = name.ends_with(".mlp.experts.gate_up_proj");
+    let is_expert_down = name.ends_with(".mlp.experts.down_proj");
+    if !(is_shared_gate_up || is_shared_down || is_expert_gate_up || is_expert_down) {
+        return Ok(None);
+    }
+    if tp_size <= 1 || tp_rank >= tp_size {
+        bail!("invalid MoE MLP TP shard: tp_rank={tp_rank}, tp_size={tp_size}");
+    }
+    let tp_size_i64 = tp_size as i64;
+    let rank = tp_rank as i64;
+    let shape = tensor.size();
+
+    if is_expert_gate_up {
+        if shape.len() != 3 || shape[1] <= 0 || shape[1] % 2 != 0 {
+            bail!(
+                "packed expert gate/up TP weight {name} must have shape [experts, 2*intermediate, hidden], got {shape:?}"
+            );
+        }
+        let intermediate = shape[1] / 2;
+        if intermediate % tp_size_i64 != 0 {
+            bail!(
+                "packed expert gate/up intermediate={intermediate} is not divisible by TP_SIZE={tp_size}"
+            );
+        }
+        let local = intermediate / tp_size_i64;
+        let gate = tensor.narrow(1, rank * local, local);
+        let up = tensor.narrow(1, intermediate + rank * local, local);
+        return Ok(Some(Tensor::cat(&[&gate, &up], 1).contiguous()));
+    }
+
+    let dim = if is_shared_gate_up {
+        if shape.len() != 2 {
+            bail!("shared expert gate/up TP weight {name} must be a matrix, got {shape:?}");
+        }
+        0
+    } else if is_shared_down {
+        if shape.len() != 2 {
+            bail!("shared expert down TP weight {name} must be a matrix, got {shape:?}");
+        }
+        1
+    } else {
+        if shape.len() != 3 {
+            bail!("routed expert down TP weight {name} must be rank 3, got {shape:?}");
+        }
+        2
+    };
+    let full = shape[dim];
+    if full <= 0 || full % tp_size_i64 != 0 {
+        bail!(
+            "MoE MLP TP weight {name} dimension {dim}={full} is not divisible by TP_SIZE={tp_size}"
+        );
+    }
+    let local = full / tp_size_i64;
+    Ok(Some(
+        tensor.narrow(dim as i64, rank * local, local).contiguous(),
+    ))
+}
+
+/// Return the rank-local frozen full-attention shard for TP. Q/K/V own
+/// contiguous output-head bundles; O owns the matching input columns.
+pub fn shard_full_attention_weight_for_tp(
+    name: &str,
+    tensor: &Tensor,
+    tp_size: usize,
+    tp_rank: usize,
+) -> Result<Option<Tensor>> {
+    let dim = if name.ends_with(".self_attn.q_proj.weight")
+        || name.ends_with(".self_attn.k_proj.weight")
+        || name.ends_with(".self_attn.v_proj.weight")
+    {
+        0
+    } else if name.ends_with(".self_attn.o_proj.weight") {
+        1
+    } else {
+        return Ok(None);
+    };
+    if tp_size <= 1 || tp_rank >= tp_size {
+        bail!("invalid full-attention TP shard: tp_rank={tp_rank}, tp_size={tp_size}");
+    }
+    let full = *tensor
+        .size()
+        .get(dim as usize)
+        .ok_or_else(|| anyhow::anyhow!("base TP attention weight {name} has no dimension {dim}"))?;
+    if full <= 0 || full % tp_size as i64 != 0 {
+        bail!(
+            "base TP attention weight {name} dimension {dim}={full} is not divisible by TP_SIZE={tp_size}"
+        );
+    }
+    let shard = full / tp_size as i64;
+    Ok(Some(
+        tensor
+            .narrow(dim, tp_rank as i64 * shard, shard)
+            .contiguous(),
+    ))
+}
+
+/// Return the rank-local frozen GDN shard for TP. Q/K/V use the model's flat
+/// `[Q_all | K_all | V_all]` layout, so QKV and depthwise-conv tensors must be
+/// sliced segment by segment before being packed into the local flat layout.
+#[allow(clippy::too_many_arguments)]
+pub fn shard_linear_attention_weight_for_tp(
+    name: &str,
+    tensor: &Tensor,
+    tp_size: usize,
+    tp_rank: usize,
+    num_k_heads: i64,
+    key_dim: i64,
+    num_v_heads: i64,
+    val_dim: i64,
+) -> Result<Option<Tensor>> {
+    let is_qkv = name.ends_with(".linear_attn.in_proj_qkv.weight");
+    let is_conv = name.ends_with(".linear_attn.conv1d.weight");
+    let is_z = name.ends_with(".linear_attn.in_proj_z.weight");
+    let is_ab = name.ends_with(".linear_attn.in_proj_a.weight")
+        || name.ends_with(".linear_attn.in_proj_b.weight");
+    let is_head = name.ends_with(".linear_attn.A_log") || name.ends_with(".linear_attn.dt_bias");
+    let is_out = name.ends_with(".linear_attn.out_proj.weight");
+    if !(is_qkv || is_conv || is_z || is_ab || is_head || is_out) {
+        return Ok(None);
+    }
+    if tp_size <= 1 || tp_rank >= tp_size {
+        bail!("invalid linear-attention TP shard: tp_rank={tp_rank}, tp_size={tp_size}");
+    }
+    let tp_size_i64 = tp_size as i64;
+    if num_k_heads <= 0
+        || num_v_heads <= 0
+        || key_dim <= 0
+        || val_dim <= 0
+        || num_v_heads % num_k_heads != 0
+        || num_k_heads % tp_size_i64 != 0
+        || num_v_heads % tp_size_i64 != 0
+    {
+        bail!(
+            "linear-attention heads/dimensions must preserve value-head groups and be divisible by TP_SIZE={tp_size}: k_heads={num_k_heads}, v_heads={num_v_heads}, key_dim={key_dim}, val_dim={val_dim}"
+        );
+    }
+
+    let q_total = num_k_heads * key_dim;
+    let v_total = num_v_heads * val_dim;
+    let qkv_total = q_total * 2 + v_total;
+    let local_q = q_total / tp_size_i64;
+    let local_v = v_total / tp_size_i64;
+    let local_heads = num_v_heads / tp_size_i64;
+    let rank = tp_rank as i64;
+    let shape = tensor.size();
+
+    let require_axis = |axis: usize, expected: i64| -> Result<()> {
+        let actual = shape.get(axis).copied().ok_or_else(|| {
+            anyhow::anyhow!("base TP linear-attention weight {name} has no dimension {axis}")
+        })?;
+        if actual != expected {
+            bail!(
+                "base TP linear-attention weight {name} dimension {axis}={actual}, expected {expected}"
+            );
+        }
+        Ok(())
+    };
+
+    if is_qkv || is_conv {
+        require_axis(0, qkv_total)?;
+        if is_qkv && shape.len() != 2 {
+            bail!("base TP linear-attention QKV weight {name} must be rank 2");
+        }
+        if is_conv && (shape.len() != 3 || shape[1] != 1) {
+            bail!(
+                "base TP linear-attention depthwise conv weight {name} must have shape [channels, 1, kernel]"
+            );
+        }
+        let q = tensor.narrow(0, rank * local_q, local_q);
+        let k = tensor.narrow(0, q_total + rank * local_q, local_q);
+        let v = tensor.narrow(0, q_total * 2 + rank * local_v, local_v);
+        return Ok(Some(Tensor::cat(&[&q, &k, &v], 0).contiguous()));
+    }
+    if is_z {
+        require_axis(0, v_total)?;
+        return Ok(Some(tensor.narrow(0, rank * local_v, local_v).contiguous()));
+    }
+    if is_ab || is_head {
+        require_axis(0, num_v_heads)?;
+        return Ok(Some(
+            tensor
+                .narrow(0, rank * local_heads, local_heads)
+                .contiguous(),
+        ));
+    }
+    require_axis(1, v_total)?;
+    Ok(Some(tensor.narrow(1, rank * local_v, local_v).contiguous()))
+}
+
 pub fn build_weight_ptrs(
     weights: &std::collections::BTreeMap<String, Tensor>,
     config: &crate::config::Qwen36RuntimeConfig,
+) -> Vec<*mut c_void> {
+    let stage = PipelineStageLayout::full(config.num_hidden_layers)
+        .expect("a full-model pipeline layout is always valid");
+    build_weight_ptrs_for_stage(weights, config, &stage)
+}
+
+pub fn build_weight_ptrs_for_stage(
+    weights: &std::collections::BTreeMap<String, Tensor>,
+    config: &crate::config::Qwen36RuntimeConfig,
+    stage: &PipelineStageLayout,
 ) -> Vec<*mut c_void> {
     let p = &config.weight_prefix;
     let mut ptrs = Vec::new();
-    for layer in 0..config.num_hidden_layers {
+    for layer in stage.layer_range.clone() {
         let lp = format!("{p}layers.{layer}");
         ptrs.push(get_ptr(weights, &format!("{lp}.input_layernorm.weight")));
-        ptrs.push(get_ptr(weights, &format!("{lp}.post_attention_layernorm.weight")));
+        ptrs.push(get_ptr(
+            weights,
+            &format!("{lp}.post_attention_layernorm.weight"),
+        ));
         match config.layer_types[layer] {
             crate::config::LayerType::FullAttention => {
                 for w in &["q_proj", "q_norm", "k_proj", "k_norm", "v_proj", "o_proj"] {
@@ -186,23 +921,50 @@ pub fn build_weight_ptrs(
                 }
             }
             crate::config::LayerType::LinearAttention => {
-                ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.in_proj_qkv.weight")));
-                ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.in_proj_z.weight")));
-                ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.in_proj_a.weight")));
-                ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.in_proj_b.weight")));
+                ptrs.push(get_ptr(
+                    weights,
+                    &format!("{lp}.linear_attn.in_proj_qkv.weight"),
+                ));
+                ptrs.push(get_ptr(
+                    weights,
+                    &format!("{lp}.linear_attn.in_proj_z.weight"),
+                ));
+                ptrs.push(get_ptr(
+                    weights,
+                    &format!("{lp}.linear_attn.in_proj_a.weight"),
+                ));
+                ptrs.push(get_ptr(
+                    weights,
+                    &format!("{lp}.linear_attn.in_proj_b.weight"),
+                ));
                 ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.A_log")));
                 ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.dt_bias")));
                 ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.conv1d.weight")));
                 ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.norm.weight")));
-                ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.out_proj.weight")));
+                ptrs.push(get_ptr(
+                    weights,
+                    &format!("{lp}.linear_attn.out_proj.weight"),
+                ));
             }
         }
         if config.is_moe {
             ptrs.push(get_ptr(weights, &format!("{lp}.mlp.gate.weight")));
-            ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert_gate.weight")));
-            ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.gate_proj.weight")));
-            ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.up_proj.weight")));
-            ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.down_proj.weight")));
+            ptrs.push(get_ptr(
+                weights,
+                &format!("{lp}.mlp.shared_expert_gate.weight"),
+            ));
+            ptrs.push(get_ptr(
+                weights,
+                &format!("{lp}.mlp.shared_expert.gate_proj.weight"),
+            ));
+            ptrs.push(get_ptr(
+                weights,
+                &format!("{lp}.mlp.shared_expert.up_proj.weight"),
+            ));
+            ptrs.push(get_ptr(
+                weights,
+                &format!("{lp}.mlp.shared_expert.down_proj.weight"),
+            ));
             ptrs.push(get_ptr(weights, &format!("{lp}.mlp.experts.gate_up_proj")));
             ptrs.push(get_ptr(weights, &format!("{lp}.mlp.experts.down_proj")));
         } else {
@@ -219,32 +981,50 @@ pub fn build_layer_configs(
     expert_start: usize,
     expert_count: usize,
 ) -> Vec<CppLayerConfig> {
-    (0..config.num_hidden_layers).map(|layer| {
-        let lt = &config.layer_types[layer];
-        CppLayerConfig {
-            layer_type: match lt { crate::config::LayerType::FullAttention => 0, _ => 1 },
-            num_heads: config.num_attention_heads,
-            num_kv_heads: config.num_key_value_heads,
-            head_dim: config.head_dim,
-            num_k_heads: config.linear_num_key_heads,
-            key_dim: config.linear_key_head_dim,
-            num_v_heads: config.linear_num_value_heads,
-            val_dim: config.linear_value_head_dim,
-            conv_kernel: config.linear_conv_kernel_dim,
-            partial_rotary_factor: config.partial_rotary_factor,
-            rope_theta: config.rope_theta,
-            rms_eps: config.rms_norm_eps,
-            num_experts: config.num_experts as i64,
-            top_k: config.num_experts_per_tok as i64,
-            moe_intermediate: config.moe_intermediate_size,
-            norm_topk_prob: if config.norm_topk_prob { 1 } else { 0 },
-            expert_start: expert_start as i64,
-            expert_count: expert_count as i64,
-            intermediate_size: config.intermediate_size,
-            nccl_comm: std::ptr::null_mut(),
-            nccl_stream: std::ptr::null_mut(),
-        }
-    }).collect()
+    let stage = PipelineStageLayout::full(config.num_hidden_layers)
+        .expect("a full-model pipeline layout is always valid");
+    build_layer_configs_for_stage(config, expert_start, expert_count, &stage)
+}
+
+pub fn build_layer_configs_for_stage(
+    config: &crate::config::Qwen36RuntimeConfig,
+    expert_start: usize,
+    expert_count: usize,
+    stage: &PipelineStageLayout,
+) -> Vec<CppLayerConfig> {
+    stage
+        .layer_range
+        .clone()
+        .map(|layer| {
+            let lt = &config.layer_types[layer];
+            CppLayerConfig {
+                layer_type: match lt {
+                    crate::config::LayerType::FullAttention => 0,
+                    _ => 1,
+                },
+                num_heads: config.num_attention_heads,
+                num_kv_heads: config.num_key_value_heads,
+                head_dim: config.head_dim,
+                num_k_heads: config.linear_num_key_heads,
+                key_dim: config.linear_key_head_dim,
+                num_v_heads: config.linear_num_value_heads,
+                val_dim: config.linear_value_head_dim,
+                conv_kernel: config.linear_conv_kernel_dim,
+                partial_rotary_factor: config.partial_rotary_factor,
+                rope_theta: config.rope_theta,
+                rms_eps: config.rms_norm_eps,
+                num_experts: config.num_experts as i64,
+                top_k: config.num_experts_per_tok as i64,
+                moe_intermediate: config.moe_intermediate_size,
+                norm_topk_prob: if config.norm_topk_prob { 1 } else { 0 },
+                expert_start: expert_start as i64,
+                expert_count: expert_count as i64,
+                intermediate_size: config.intermediate_size,
+                nccl_comm: std::ptr::null_mut(),
+                nccl_stream: std::ptr::null_mut(),
+            }
+        })
+        .collect()
 }
 
 /// Build weight pointers for MTP layers (full attention layers).
@@ -258,7 +1038,10 @@ pub fn build_mtp_weight_ptrs(
     for layer in 0..config.mtp_num_hidden_layers {
         let lp = format!("mtp.layers.{layer}");
         ptrs.push(get_ptr(weights, &format!("{lp}.input_layernorm.weight")));
-        ptrs.push(get_ptr(weights, &format!("{lp}.post_attention_layernorm.weight")));
+        ptrs.push(get_ptr(
+            weights,
+            &format!("{lp}.post_attention_layernorm.weight"),
+        ));
         // Full attention: q, q_norm, k, k_norm, v, o
         for w in &["q_proj", "q_norm", "k_proj", "k_norm", "v_proj", "o_proj"] {
             ptrs.push(get_ptr(weights, &format!("{lp}.self_attn.{w}.weight")));
@@ -266,10 +1049,22 @@ pub fn build_mtp_weight_ptrs(
         if config.is_moe {
             // MoE: gate, shared_expert_gate, shared_gate_proj, shared_up_proj, shared_down_proj, experts_gate_up, experts_down
             ptrs.push(get_ptr(weights, &format!("{lp}.mlp.gate.weight")));
-            ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert_gate.weight")));
-            ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.gate_proj.weight")));
-            ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.up_proj.weight")));
-            ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.down_proj.weight")));
+            ptrs.push(get_ptr(
+                weights,
+                &format!("{lp}.mlp.shared_expert_gate.weight"),
+            ));
+            ptrs.push(get_ptr(
+                weights,
+                &format!("{lp}.mlp.shared_expert.gate_proj.weight"),
+            ));
+            ptrs.push(get_ptr(
+                weights,
+                &format!("{lp}.mlp.shared_expert.up_proj.weight"),
+            ));
+            ptrs.push(get_ptr(
+                weights,
+                &format!("{lp}.mlp.shared_expert.down_proj.weight"),
+            ));
             ptrs.push(get_ptr(weights, &format!("{lp}.mlp.experts.gate_up_proj")));
             ptrs.push(get_ptr(weights, &format!("{lp}.mlp.experts.down_proj")));
         } else {
@@ -288,31 +1083,33 @@ pub fn build_mtp_layer_configs(
     expert_start: usize,
     expert_count: usize,
 ) -> Vec<CppLayerConfig> {
-    (0..config.mtp_num_hidden_layers).map(|_| {
-        CppLayerConfig {
-            layer_type: 0,  // MTP layers are always full attention
-            num_heads: config.num_attention_heads,
-            num_kv_heads: config.num_key_value_heads,
-            head_dim: config.head_dim,
-            num_k_heads: config.linear_num_key_heads,
-            key_dim: config.linear_key_head_dim,
-            num_v_heads: config.linear_num_value_heads,
-            val_dim: config.linear_value_head_dim,
-            conv_kernel: config.linear_conv_kernel_dim,
-            partial_rotary_factor: config.partial_rotary_factor,
-            rope_theta: config.rope_theta,
-            rms_eps: config.rms_norm_eps,
-            num_experts: config.num_experts as i64,
-            top_k: config.num_experts_per_tok as i64,
-            moe_intermediate: config.moe_intermediate_size,
-            norm_topk_prob: if config.norm_topk_prob { 1 } else { 0 },
-            expert_start: expert_start as i64,
-            expert_count: expert_count as i64,
-            intermediate_size: config.intermediate_size,
-            nccl_comm: std::ptr::null_mut(),
-            nccl_stream: std::ptr::null_mut(),
-        }
-    }).collect()
+    (0..config.mtp_num_hidden_layers)
+        .map(|_| {
+            CppLayerConfig {
+                layer_type: 0, // MTP layers are always full attention
+                num_heads: config.num_attention_heads,
+                num_kv_heads: config.num_key_value_heads,
+                head_dim: config.head_dim,
+                num_k_heads: config.linear_num_key_heads,
+                key_dim: config.linear_key_head_dim,
+                num_v_heads: config.linear_num_value_heads,
+                val_dim: config.linear_value_head_dim,
+                conv_kernel: config.linear_conv_kernel_dim,
+                partial_rotary_factor: config.partial_rotary_factor,
+                rope_theta: config.rope_theta,
+                rms_eps: config.rms_norm_eps,
+                num_experts: config.num_experts as i64,
+                top_k: config.num_experts_per_tok as i64,
+                moe_intermediate: config.moe_intermediate_size,
+                norm_topk_prob: if config.norm_topk_prob { 1 } else { 0 },
+                expert_start: expert_start as i64,
+                expert_count: expert_count as i64,
+                intermediate_size: config.intermediate_size,
+                nccl_comm: std::ptr::null_mut(),
+                nccl_stream: std::ptr::null_mut(),
+            }
+        })
+        .collect()
 }
 
 /// Opaque training context handle.
@@ -322,72 +1119,243 @@ pub struct CppTrainingContext {
 }
 
 impl CppTrainingContext {
+    /// Whether the loaded native library supports complete tenant-local Adam
+    /// overrides (beta1, beta2, and epsilon) in addition to learning rate.
+    pub fn supports_complete_dynamic_adam(&self) -> bool {
+        get_kernels().is_some_and(|kernels| {
+            kernels.add_lora_with_optimizer_v2.is_some()
+                && kernels.add_lora_for_restore_with_optimizer_v2.is_some()
+        })
+    }
+
     /// Create training context — LoRA A/B created in C++ as at::Tensor (requires_grad=true).
     pub fn new(
         weights: &std::collections::BTreeMap<String, Tensor>,
         config: &crate::config::Qwen36RuntimeConfig,
         compute_kind: Kind,
-        lr: f64, beta1: f64, beta2: f64, eps: f64,
+        lr: f64,
+        beta1: f64,
+        beta2: f64,
+        eps: f64,
+        lora_scaling: f64,
+        lora_rank: i64,
+        base_tp_attention: bool,
+        base_tp_mlp: bool,
+        vocab_parallel: bool,
+        data_parallel: bool,
+        expert_parallel: bool,
+        target_layers: &[usize],
+        target_modules: &[Qwen36LoraTargetModule],
+        expert_start: usize,
+        expert_count: usize,
+    ) -> Result<Self> {
+        let stage = PipelineStageLayout::full(config.num_hidden_layers)?;
+        Self::new_for_stage(
+            weights,
+            config,
+            &stage,
+            compute_kind,
+            lr,
+            beta1,
+            beta2,
+            eps,
+            lora_scaling,
+            lora_rank,
+            base_tp_attention,
+            base_tp_mlp,
+            vocab_parallel,
+            data_parallel,
+            expert_parallel,
+            target_layers,
+            target_modules,
+            expert_start,
+            expert_count,
+        )
+    }
+
+    /// Create a context that owns exactly one contiguous pipeline stage.
+    /// Target layers remain global IDs and are mapped to local slots by C++.
+    #[allow(clippy::too_many_arguments)]
+    pub fn new_for_stage(
+        weights: &std::collections::BTreeMap<String, Tensor>,
+        config: &crate::config::Qwen36RuntimeConfig,
+        stage: &PipelineStageLayout,
+        compute_kind: Kind,
+        lr: f64,
+        beta1: f64,
+        beta2: f64,
+        eps: f64,
         lora_scaling: f64,
         lora_rank: i64,
+        base_tp_attention: bool,
+        base_tp_mlp: bool,
+        vocab_parallel: bool,
+        data_parallel: bool,
+        expert_parallel: bool,
         target_layers: &[usize],
+        target_modules: &[Qwen36LoraTargetModule],
         expert_start: usize,
         expert_count: usize,
     ) -> Result<Self> {
+        if stage.global_num_layers != config.num_hidden_layers {
+            bail!(
+                "pipeline layout has {} global layers, but model config has {}",
+                stage.global_num_layers,
+                config.num_hidden_layers
+            );
+        }
         let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
-        let mut weight_ptrs = build_weight_ptrs(weights, config);
-        let layer_configs = build_layer_configs(config, expert_start, expert_count);
+        let weight_ptrs = build_weight_ptrs_for_stage(weights, config, stage);
+        let layer_configs =
+            build_layer_configs_for_stage(config, expert_start, expert_count, stage);
 
-        let embed_ptr = get_ptr(weights, &format!("{}embed_tokens.weight", config.weight_prefix));
-        let final_norm_ptr = get_ptr(weights, &format!("{}norm.weight", config.weight_prefix));
-        let lm_head_ptr = if config.tie_word_embeddings {
-            // Tied embeddings: use embed_tokens as lm_head
-            embed_ptr
+        let embedding_weight_ptr = get_ptr(
+            weights,
+            &format!("{}embed_tokens.weight", config.weight_prefix),
+        );
+        let embed_ptr = if stage.is_first() {
+            embedding_weight_ptr
+        } else {
+            std::ptr::null_mut()
+        };
+        let final_norm_ptr = if stage.is_last() {
+            get_ptr(weights, &format!("{}norm.weight", config.weight_prefix))
+        } else {
+            std::ptr::null_mut()
+        };
+        let lm_head_ptr = if !stage.is_last() {
+            std::ptr::null_mut()
+        } else if config.tie_word_embeddings {
+            // Tied models duplicate the frozen vocabulary shard on the two
+            // boundary stages; the last stage receives it only as LM head.
+            embedding_weight_ptr
         } else {
             get_ptr(weights, "lm_head.weight")
         };
 
         let compute_type = compute_kind as i32;
 
-        // LEAK the Vecs so their backing arrays survive (C++ holds raw pointers)
+        // The C++ constructor copies these pointer/config arrays into its own
+        // vectors before returning. The tensors they point to remain owned by
+        // the Rust session for the full context lifetime.
         let wp_ptr = weight_ptrs.as_ptr() as *mut *mut c_void;
         let wp_len = weight_ptrs.len();
-        std::mem::forget(weight_ptrs);
         let lc_ptr = layer_configs.as_ptr() as *mut c_void;
-        std::mem::forget(layer_configs);
 
-        // Convert target_layers to i64 array (leaked to keep alive)
+        // The constructor also consumes the target layer array synchronously.
         let target_i64: Vec<i64> = target_layers.iter().map(|&x| x as i64).collect();
         let tl_ptr = if target_i64.is_empty() {
             std::ptr::null()
         } else {
-            let p = target_i64.as_ptr();
-            std::mem::forget(target_i64);
-            p
+            target_i64.as_ptr()
         };
         let tl_len = target_layers.len() as i64;
 
+        let module_names = target_modules
+            .iter()
+            .map(Qwen36LoraTargetModule::cpp_name)
+            .collect::<Vec<_>>()
+            .join(",");
+        let module_names_c = std::ffi::CString::new(module_names)
+            .map_err(|_| anyhow::anyhow!("LoRA target module contains NUL"))?;
+        let modules_ptr = if target_modules.is_empty() {
+            std::ptr::null()
+        } else {
+            module_names_c.as_ptr()
+        };
+
         let ptr = unsafe {
             (kh.create_ctx)(
-                wp_ptr, wp_len as i64,
-                embed_ptr, final_norm_ptr, lm_head_ptr,
-                lc_ptr, config.num_hidden_layers as i64,
-                compute_type, lora_scaling,
-                lr, beta1, beta2, eps,
-                config.vocab_size, config.rms_norm_eps,
-                lora_rank, tl_ptr, tl_len,
+                wp_ptr,
+                wp_len as i64,
+                embed_ptr,
+                final_norm_ptr,
+                lm_head_ptr,
+                lc_ptr,
+                stage.local_num_layers() as i64,
+                stage.layer_range.start as i64,
+                stage.global_num_layers as i64,
+                stage.native_flags(),
+                compute_type,
+                lora_scaling,
+                lr,
+                beta1,
+                beta2,
+                eps,
+                config.vocab_size,
+                config.rms_norm_eps,
+                lora_rank,
+                tl_ptr,
+                tl_len,
+                modules_ptr,
+                i32::from(base_tp_attention)
+                    | (i32::from(data_parallel) << 1)
+                    | (i32::from(vocab_parallel) << 2)
+                    | (i32::from(expert_parallel) << 3)
+                    | (i32::from(base_tp_mlp) << 4),
             )
         };
         if ptr.is_null() {
             bail!("C++ create_training_context returned null");
         }
+        let router_aux_status =
+            unsafe { (kh.set_router_aux_loss_coef)(ptr, config.router_aux_loss_coef) };
+        if router_aux_status != 0 {
+            unsafe { (kh.free_ctx)(ptr) };
+            bail!(
+                "C++ router auxiliary loss configuration failed for coefficient {}",
+                config.router_aux_loss_coef
+            );
+        }
+        let base_tp_status = unsafe { (kh.set_base_tp_mlp)(ptr, i32::from(base_tp_mlp)) };
+        if base_tp_status != 0 {
+            unsafe { (kh.free_ctx)(ptr) };
+            bail!("C++ base dense MLP TP configuration failed");
+        }
         let lora_count = unsafe { (kh.get_lora_count)(ptr) };
         Ok(Self { ptr, lora_count })
     }
 
+    /// Configure native Megatron-style global gradient clipping. Zero disables
+    /// clipping; dynamic multi-LoRA applies the threshold independently per tenant.
+    pub fn set_max_grad_norm(&self, max_grad_norm: f64) -> Result<()> {
+        if !max_grad_norm.is_finite() || max_grad_norm < 0.0 {
+            bail!("native max_grad_norm must be finite and non-negative");
+        }
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let status = unsafe { (kh.set_max_grad_norm)(self.ptr, max_grad_norm) };
+        if status != 0 {
+            bail!("C++ max gradient norm configuration failed");
+        }
+        Ok(())
+    }
+
+    /// Configure the native padding token so training can derive its
+    /// attention mask inside the C++ kernel. This keeps GPU mask construction
+    /// out of the Rust training loop.
+    pub fn set_pad_token_id(&self, pad_token_id: i64) -> Result<()> {
+        if pad_token_id < 0 {
+            bail!("native pad_token_id must be non-negative");
+        }
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let setter = kh
+            .set_pad_token_id
+            .ok_or_else(|| anyhow::anyhow!("native kernel lacks pad-token mask support"))?;
+        let status = unsafe { setter(self.ptr, pad_token_id) };
+        if status != 0 {
+            bail!("C++ pad_token_id configuration failed");
+        }
+        Ok(())
+    }
+
     /// Run one training step: forward + loss + backward + Adam update.
     /// Returns loss value.
-    pub fn train_step(&self, input_ids: &Tensor, target_mask: &Tensor, attention_mask: &Tensor) -> Result<f64> {
+    pub fn train_step(
+        &self,
+        input_ids: &Tensor,
+        target_mask: &Tensor,
+        attention_mask: &Tensor,
+    ) -> Result<f64> {
         let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
         let loss = unsafe {
             (kh.train_step)(
@@ -397,172 +1365,1359 @@ impl CppTrainingContext {
                 attention_mask.as_ptr() as *mut _,
             )
         };
-        if loss < 0.0 {
+        if !loss.is_finite() || loss < 0.0 {
             bail!("C++ train_step failed");
         }
         Ok(loss)
     }
 
-    /// Train all adapters in batched chunks. Each chunk runs independent
-    /// forward → loss → backward → Adam. Input is expanded to [N, seq].
-    /// n_total: total number of adapters. lora_rank: LoRA rank for N_max calc.
-    /// Returns average loss across chunks.
-    pub fn train_multi_lora(
-        &self, input_ids: &Tensor, target_mask: &Tensor, attention_mask: &Tensor,
-        n_total: i32, lora_rank: i32,
+    /// Run one micro-batch and optionally apply the synchronized Adam update.
+    pub fn train_micro_step(
+        &self,
+        input_ids: &Tensor,
+        target_mask: &Tensor,
+        attention_mask: Option<&Tensor>,
+        gradient_scale: f64,
+        apply_optimizer: bool,
     ) -> Result<f64> {
         let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
         let loss = unsafe {
-            (kh.train_multi_lora)(
+            (kh.train_micro_step)(
                 self.ptr,
                 input_ids.as_ptr() as *mut _,
                 target_mask.as_ptr() as *mut _,
-                attention_mask.as_ptr() as *mut _,
-                n_total,
-                lora_rank,
+                attention_mask
+                    .map_or(std::ptr::null_mut(), |tensor| tensor.as_ptr() as *mut _),
+                gradient_scale,
+                i32::from(apply_optimizer),
             )
         };
-        if loss < 0.0 {
-            bail!("C++ train_multi_lora failed");
+        if !loss.is_finite() || loss < 0.0 {
+            bail!("C++ train_micro_step failed");
         }
         Ok(loss)
     }
 
-    /// Get LoRA A tensor by index (for saving).
-    pub fn get_lora_a(&self, index: i64) -> Option<Tensor> {
-        let kh = get_kernels()?;
-        let ptr = unsafe { (kh.get_lora_a)(self.ptr, index) };
-        if ptr.is_null() { return None; }
-        Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) })
-    }
-
-    /// Get LoRA B tensor by index (for saving).
-    pub fn get_lora_b(&self, index: i64) -> Option<Tensor> {
-        let kh = get_kernels()?;
-        let ptr = unsafe { (kh.get_lora_b)(self.ptr, index) };
-        if ptr.is_null() { return None; }
-        Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) })
+    /// Open a fixed-shape non-interleaved 1F1B pipeline window. The native side owns all
+    /// activation/gradient slots until `pipeline_finish_v1` or abort.
+    pub fn pipeline_begin_v1(&self, window_id: i64, num_microbatches: i64) -> Result<()> {
+        self.pipeline_begin_with_flags_v1(window_id, num_microbatches, 0)
     }
 
-    pub fn lora_count(&self) -> i64 {
-        self.lora_count
+    /// Open a dynamic multi-tenant LoRA 1F1B window in native registry order.
+    pub fn pipeline_begin_dynamic_v1(
+        &self,
+        window_id: i64,
+        num_microbatches: i64,
+    ) -> Result<()> {
+        self.pipeline_begin_with_flags_v1(
+            window_id,
+            num_microbatches,
+            PIPELINE_WINDOW_FLAG_DYNAMIC_LORA,
+        )
     }
 
-    /// Set MTP weights on the C++ training context.
-    /// Must be called after `new()` if MTP is enabled.
-    pub fn set_mtp_weights(
+    /// Open a dynamic pipeline window for only the requested live tenants.
+    /// Unselected adapters retain their parameters, accumulators, and clocks.
+    pub fn pipeline_begin_dynamic_selected_v1(
         &self,
-        weights: &std::collections::BTreeMap<String, Tensor>,
-        config: &crate::config::Qwen36RuntimeConfig,
-        expert_start: usize,
-        expert_count: usize,
+        window_id: i64,
+        num_microbatches: i64,
+        adapter_ids: &[i64],
     ) -> Result<()> {
+        if adapter_ids.is_empty() {
+            bail!("selected dynamic pipeline requires at least one adapter ID");
+        }
+        if window_id < 0 || num_microbatches <= 0 {
+            bail!("pipeline window id and microbatch count must be positive");
+        }
         let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
-
-        let mtp_fc_ptr = get_ptr(weights, "mtp.fc.weight");
-        let mtp_pre_fc_norm_emb_ptr = get_ptr(weights, "mtp.pre_fc_norm_embedding.weight");
-        let mtp_pre_fc_norm_hidden_ptr = get_ptr(weights, "mtp.pre_fc_norm_hidden.weight");
-        let mtp_norm_ptr = get_ptr(weights, "mtp.norm.weight");
-
-        let mut mtp_weight_ptrs = build_mtp_weight_ptrs(weights, config);
-        let mtp_layer_configs = build_mtp_layer_configs(config, expert_start, expert_count);
-
-        let wp_ptr = mtp_weight_ptrs.as_ptr() as *mut *mut c_void;
-        let wp_len = mtp_weight_ptrs.len();
-        std::mem::forget(mtp_weight_ptrs);
-        let lc_ptr = mtp_layer_configs.as_ptr() as *mut c_void;
-        std::mem::forget(mtp_layer_configs);
-
-        unsafe {
-            (kh.set_mtp_weights)(
-                self.ptr,
-                mtp_fc_ptr,
-                mtp_pre_fc_norm_emb_ptr,
-                mtp_pre_fc_norm_hidden_ptr,
-                mtp_norm_ptr,
-                wp_ptr,
-                wp_len as i64,
-                lc_ptr,
-                config.mtp_num_hidden_layers as i64,
-            );
+        let begin = kh.pipeline_begin_selected.ok_or_else(|| {
+            anyhow::anyhow!("native kernels do not expose selected dynamic pipeline ABI")
+        })?;
+        let spec = PipelineWindowV1 {
+            struct_size: std::mem::size_of::<PipelineWindowV1>() as u32,
+            version: 1,
+            window_id,
+            num_microbatches,
+            schedule: 0,
+            num_chunks: 1,
+            flags: PIPELINE_WINDOW_FLAG_DYNAMIC_LORA,
+        };
+        let count = i32::try_from(adapter_ids.len())
+            .context("selected adapter count exceeds i32")?;
+        let status = unsafe { begin(self.ptr, &spec, adapter_ids.as_ptr(), count) };
+        if status != 0 {
+            bail!("C++ selected dynamic pipeline begin failed with status {status}");
         }
         Ok(())
     }
 
-    /// Enable/disable gradient checkpointing.
-    /// When enabled, layers are grouped by `group_size` and intermediate
-    /// activations are recomputed during backward instead of stored.
-    pub fn set_checkpoint(&self, enable: bool, group_size: i64) {
-        let kh = get_kernels().expect("kernels not loaded");
-        unsafe {
-            (kh.set_checkpoint)(self.ptr, if enable { 1 } else { 0 }, group_size);
+    fn pipeline_begin_with_flags_v1(
+        &self,
+        window_id: i64,
+        num_microbatches: i64,
+        flags: i32,
+    ) -> Result<()> {
+        if window_id < 0 {
+            bail!("pipeline window id must be non-negative");
         }
-    }
-
-    /// Set NCCL communicator for Expert Parallel all-reduce.
-    /// Must be called after `new()` if EP is enabled.
-    /// comm_ptr / stream_ptr from `NcclPersistentComm::raw_comm_ptr()` / `raw_stream_ptr()`.
-    pub fn set_nccl_comm(&self, comm_ptr: *mut c_void, stream_ptr: *mut c_void, ep_rank: i32, ep_world_size: i32) {
-        let kh = get_kernels().expect("kernels not loaded");
-        unsafe {
-            (kh.set_nccl_comm)(self.ptr, comm_ptr, stream_ptr, ep_rank, ep_world_size);
+        if num_microbatches <= 0 {
+            bail!("pipeline window must contain at least one microbatch");
         }
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let spec = PipelineWindowV1 {
+            struct_size: std::mem::size_of::<PipelineWindowV1>() as u32,
+            version: 1,
+            window_id,
+            num_microbatches,
+            schedule: 0,
+            num_chunks: 1,
+            flags,
+        };
+        let status = unsafe { (kh.pipeline_begin)(self.ptr, &spec) };
+        if status != 0 {
+            bail!("C++ pipeline window begin failed with status {status}");
+        }
+        Ok(())
     }
 
-    /// Initialize NCCL communicator directly in C++ (preferred over set_nccl_comm).
-    /// Reads RANK/WORLD_SIZE/LOCAL_RANK from env vars.
-    /// Returns 0 on success, -1 on failure.
-    pub fn init_nccl(&self) -> i32 {
-        let kh = get_kernels().expect("kernels not loaded");
-        unsafe { (kh.init_nccl)(self.ptr) }
+    /// Advance a pipeline window by one canonical non-interleaved 1F1B tick.
+    pub fn pipeline_tick_v1(
+        &self,
+        window_id: i64,
+        forward_mb: Option<i64>,
+        backward_mb: Option<i64>,
+        input_ids: Option<&Tensor>,
+        target_mask: Option<&Tensor>,
+        attention_mask: Option<&Tensor>,
+        gradient_scale: f64,
+    ) -> Result<PipelineResultV1> {
+        if !gradient_scale.is_finite() || gradient_scale <= 0.0 {
+            bail!("pipeline gradient scale must be finite and positive");
+        }
+        if forward_mb.is_some() != input_ids.is_some()
+            || forward_mb.is_some() != target_mask.is_some()
+            || forward_mb.is_some() != attention_mask.is_some()
+        {
+            bail!("pipeline forward microbatch requires all input tensors");
+        }
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let tick = PipelineTickV1 {
+            struct_size: std::mem::size_of::<PipelineTickV1>() as u32,
+            version: 1,
+            window_id,
+            forward_mb: forward_mb.unwrap_or(-1),
+            backward_mb: backward_mb.unwrap_or(-1),
+            chunk_id: 0,
+            phase: match (forward_mb, backward_mb) {
+                (Some(_), None) => 0,
+                (Some(_), Some(_)) => 1,
+                (None, _) => 2,
+            },
+            input_ids: input_ids
+                .map(|tensor| tensor.as_ptr() as *mut c_void)
+                .unwrap_or(std::ptr::null_mut()),
+            target_mask: target_mask
+                .map(|tensor| tensor.as_ptr() as *mut c_void)
+                .unwrap_or(std::ptr::null_mut()),
+            attention_mask: attention_mask
+                .map(|tensor| tensor.as_ptr() as *mut c_void)
+                .unwrap_or(std::ptr::null_mut()),
+            gradient_scale,
+        };
+        let mut result = PipelineResultV1 {
+            struct_size: std::mem::size_of::<PipelineResultV1>() as u32,
+            version: 1,
+            status: 0,
+            completed_fwd: 0,
+            completed_bwd: 0,
+            in_flight: 0,
+            optimizer_step: 0,
+            loss: 0.0,
+        };
+        let status = unsafe { (kh.pipeline_tick)(self.ptr, &tick, &mut result) };
+        if status != 0 || result.status != 0 {
+            bail!("C++ pipeline window tick failed with status {status}");
+        }
+        Ok(result)
     }
 
-    /// Set CUDA device and force PyTorch CUDA context initialization.
-    /// Must be called before any GPU operation in worker processes.
-    pub fn set_cuda_device(device: i32) {
-        let kh = get_kernels().expect("kernels not loaded");
-        unsafe { (kh.set_cuda_device)(device) }
+    /// Finish a pipeline window and apply its single accumulated optimizer step.
+    pub fn pipeline_finish_v1(&self) -> Result<PipelineResultV1> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let mut result = PipelineResultV1 {
+            struct_size: std::mem::size_of::<PipelineResultV1>() as u32,
+            version: 1,
+            status: 0,
+            completed_fwd: 0,
+            completed_bwd: 0,
+            in_flight: 0,
+            optimizer_step: 0,
+            loss: 0.0,
+        };
+        let status = unsafe { (kh.pipeline_finish)(self.ptr, 1, &mut result) };
+        if status != 0 || result.status != 0 {
+            bail!("C++ pipeline window finish failed with status {status}");
+        }
+        Ok(result)
     }
 
-    /// Add a new LoRA adapter. Returns adapter ID (>0) on success.
-    /// target_layers: empty = all layers
-    /// target_modules: comma-separated, e.g. "q_proj,k_proj,v_proj,o_proj". Empty = all.
-    pub fn add_lora(
-        &self, rank: i64, alpha: f64,
-        target_layers: &[i64], target_modules: &str,
-    ) -> Result<i64> {
+    /// Finish a selected dynamic PP window and return one loss per tenant.
+    pub fn pipeline_finish_dynamic_report_v1(
+        &self,
+        adapter_count: usize,
+    ) -> Result<(PipelineResultV1, Vec<f64>)> {
         let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
-        let tl_ptr = if target_layers.is_empty() { std::ptr::null() } else { target_layers.as_ptr() };
-        let tl_len = target_layers.len() as i64;
-        let modules_c = std::ffi::CString::new(target_modules).unwrap();
-        let modules_ptr = if target_modules.is_empty() { std::ptr::null() } else { modules_c.as_ptr() };
-        let id = unsafe {
-            (kh.add_lora)(self.ptr, rank, alpha, tl_ptr, tl_len, modules_ptr)
+        let report = kh.pipeline_finish_report.ok_or_else(|| {
+            anyhow::anyhow!("loaded kernel does not provide dynamic pipeline loss reports")
+        })?;
+        let mut result = PipelineResultV1 {
+            struct_size: std::mem::size_of::<PipelineResultV1>() as u32,
+            version: 1,
+            status: 0,
+            completed_fwd: 0,
+            completed_bwd: 0,
+            in_flight: 0,
+            optimizer_step: 0,
+            loss: 0.0,
         };
-        if id < 0 {
-            bail!("C++ add_lora failed");
+        let mut adapter_losses = vec![f64::NAN; adapter_count];
+        let status = unsafe {
+            report(
+                self.ptr,
+                1,
+                &mut result,
+                adapter_losses.as_mut_ptr(),
+                i32::try_from(adapter_count).context("adapter count exceeds i32")?,
+            )
+        };
+        if status != 0
+            || result.status != 0
+            || adapter_losses
+                .iter()
+                .any(|loss| !loss.is_finite() || *loss < 0.0)
+        {
+            bail!("C++ dynamic pipeline loss report failed with status {status}");
         }
-        Ok(id)
+        Ok((result, adapter_losses))
     }
 
-    /// Remove a LoRA adapter by ID.
-    pub fn remove_lora(&self, adapter_id: i64) -> Result<bool> {
+    /// Abort a pipeline window and discard any accumulated gradients.
+    pub fn pipeline_abort_v1(&self) -> Result<()> {
         let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
-        let found = unsafe { (kh.remove_lora)(self.ptr, adapter_id) };
-        Ok(found != 0)
+        let status = unsafe { (kh.pipeline_abort)(self.ptr) };
+        if status != 0 {
+            bail!("C++ pipeline window abort failed with status {status}");
+        }
+        Ok(())
+    }
+
+    /// Train every live dynamic adapter, including heterogeneous signatures.
+    /// The rank argument is retained for source compatibility with ABI22 callers.
+    pub fn train_multi_lora(
+        &self,
+        input_ids: &Tensor,
+        target_mask: &Tensor,
+        attention_mask: &Tensor,
+        n_total: i32,
+        lora_rank: i32,
+    ) -> Result<f64> {
+        if n_total <= 0 {
+            bail!("n_total must be positive, got {n_total}");
+        }
+        let adapter_ids = self.list_dynamic_lora();
+        if adapter_ids.len() != n_total as usize {
+            bail!(
+                "live adapter count {} does not match n_total={n_total}",
+                adapter_ids.len()
+            );
+        }
+        self.train_multi_lora_selected(
+            input_ids,
+            target_mask,
+            attention_mask,
+            &adapter_ids,
+            lora_rank,
+        )
+    }
+
+    /// Train only the requested dynamic adapters. The native wrapper scopes
+    /// its registry to these IDs so unselected tenants keep their parameters,
+    /// gradients, and optimizer clocks untouched.
+    pub fn train_multi_lora_selected(
+        &self,
+        input_ids: &Tensor,
+        target_mask: &Tensor,
+        attention_mask: &Tensor,
+        adapter_ids: &[i64],
+        _lora_rank: i32,
+    ) -> Result<f64> {
+        if adapter_ids.is_empty() {
+            bail!("selected multi-LoRA requires at least one adapter ID");
+        }
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let loss = unsafe {
+            (kh.train_multi_lora_selected_v2)(
+                self.ptr,
+                input_ids.as_ptr() as *mut _,
+                target_mask.as_ptr() as *mut _,
+                attention_mask.as_ptr() as *mut _,
+                adapter_ids.as_ptr(),
+                i32::try_from(adapter_ids.len()).context("selected adapter count exceeds i32")?,
+            )
+        };
+        if !loss.is_finite() || loss < 0.0 {
+            if !self.is_healthy() {
+                bail!("native dynamic LoRA context is poisoned; recreate the session");
+            }
+            bail!("C++ train_multi_lora_selected_v2 failed");
+        }
+        Ok(loss)
+    }
+
+    /// Train selected adapters once and return globally normalized losses in
+    /// the same order as `adapter_ids`.
+    pub fn train_multi_lora_selected_report(
+        &self,
+        input_ids: &Tensor,
+        target_mask: &Tensor,
+        attention_mask: &Tensor,
+        adapter_ids: &[i64],
+    ) -> Result<MultiLoraLossReport> {
+        if adapter_ids.is_empty() {
+            bail!("selected multi-LoRA loss report requires adapter IDs");
+        }
+        let adapter_count =
+            i32::try_from(adapter_ids.len()).context("selected adapter count exceeds i32")?;
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let report = kh.train_multi_lora_selected_v3;
+        let mut aggregate_loss = f64::NAN;
+        let mut adapter_losses = vec![f64::NAN; adapter_ids.len()];
+        let status = unsafe {
+            report(
+                self.ptr,
+                input_ids.as_ptr() as *mut _,
+                target_mask.as_ptr() as *mut _,
+                attention_mask.as_ptr() as *mut _,
+                adapter_ids.as_ptr(),
+                adapter_count,
+                &mut aggregate_loss,
+                adapter_losses.as_mut_ptr(),
+                adapter_count,
+            )
+        };
+        if status != 0
+            || !aggregate_loss.is_finite()
+            || aggregate_loss < 0.0
+            || adapter_losses
+                .iter()
+                .any(|loss| !loss.is_finite() || *loss < 0.0)
+        {
+            if !self.is_healthy() {
+                bail!("native dynamic LoRA context is poisoned; recreate the session");
+            }
+            bail!("C++ train_multi_lora_selected_v3 failed");
+        }
+        Ok(MultiLoraLossReport {
+            aggregate_loss,
+            adapter_losses,
+        })
+    }
+
+    /// Run one complete fixed-LoRA step from a borrowed host int64 batch.
+    /// The native entry owns validation, H2D copies, forward/backward, and Adam.
+    pub fn train_step_host_i64(
+        &self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+    ) -> Result<f64> {
+        validate_host_batch(input_ids, target_mask, attention_mask, batch_size, seq_len)?;
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let loss = unsafe {
+            (kh.train_step_host_i64)(
+                self.ptr,
+                input_ids.as_ptr(),
+                target_mask.as_ptr(),
+                attention_mask.as_ptr(),
+                i64::try_from(batch_size).context("host batch_size exceeds i64")?,
+                i64::try_from(seq_len).context("host seq_len exceeds i64")?,
+            )
+        };
+        if !loss.is_finite() || loss < 0.0 {
+            bail!("C++ train_step_host_i64 failed");
+        }
+        Ok(loss)
+    }
+
+    /// Run one complete dynamic multi-LoRA step from a borrowed host batch.
+    pub fn train_multi_lora_host_i64(
+        &self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+        n_total: i32,
+        lora_rank: i32,
+        adapter_ids: &[i64],
+    ) -> Result<f64> {
+        validate_host_batch(input_ids, target_mask, attention_mask, batch_size, seq_len)?;
+        if n_total <= 0 {
+            bail!("n_total must be positive, got {n_total}");
+        }
+        if !adapter_ids.is_empty() && adapter_ids.len() != n_total as usize {
+            bail!(
+                "selected adapter count {} does not match n_total={n_total}",
+                adapter_ids.len()
+            );
+        }
+        let adapter_count =
+            i32::try_from(adapter_ids.len()).context("selected adapter count exceeds i32")?;
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let loss = unsafe {
+            (kh.train_multi_lora_host_i64)(
+                self.ptr,
+                input_ids.as_ptr(),
+                target_mask.as_ptr(),
+                attention_mask.as_ptr(),
+                i64::try_from(batch_size).context("host batch_size exceeds i64")?,
+                i64::try_from(seq_len).context("host seq_len exceeds i64")?,
+                n_total,
+                lora_rank,
+                if adapter_ids.is_empty() {
+                    std::ptr::null()
+                } else {
+                    adapter_ids.as_ptr()
+                },
+                adapter_count,
+            )
+        };
+        if !loss.is_finite() || loss < 0.0 {
+            if !self.is_healthy() {
+                bail!("native dynamic LoRA context is poisoned; recreate the session");
+            }
+            bail!("C++ train_multi_lora_host_i64 failed");
+        }
+        Ok(loss)
+    }
+
+    /// Borrow a host batch for selected adapters and return one globally
+    /// normalized loss per adapter alongside the scalar entry point.
+    pub fn train_multi_lora_host_i64_report(
+        &self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+        n_total: i32,
+        lora_rank: i32,
+        adapter_ids: &[i64],
+    ) -> Result<MultiLoraLossReport> {
+        validate_host_batch(input_ids, target_mask, attention_mask, batch_size, seq_len)?;
+        if n_total <= 0 || adapter_ids.len() != n_total as usize {
+            bail!("multi-LoRA loss report requires n_total positive selected adapter IDs");
+        }
+        let adapter_count =
+            i32::try_from(adapter_ids.len()).context("selected adapter count exceeds i32")?;
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let report = kh.train_multi_lora_host_i64_v2;
+        let mut aggregate_loss = f64::NAN;
+        let mut adapter_losses = vec![f64::NAN; adapter_ids.len()];
+        let status = unsafe {
+            report(
+                self.ptr,
+                input_ids.as_ptr(),
+                target_mask.as_ptr(),
+                attention_mask.as_ptr(),
+                i64::try_from(batch_size).context("host batch_size exceeds i64")?,
+                i64::try_from(seq_len).context("host seq_len exceeds i64")?,
+                n_total,
+                lora_rank,
+                adapter_ids.as_ptr(),
+                adapter_count,
+                &mut aggregate_loss,
+                adapter_losses.as_mut_ptr(),
+                adapter_count,
+            )
+        };
+        if status != 0
+            || !aggregate_loss.is_finite()
+            || aggregate_loss < 0.0
+            || adapter_losses
+                .iter()
+                .any(|loss| !loss.is_finite() || *loss < 0.0)
+        {
+            if !self.is_healthy() {
+                bail!("native dynamic LoRA context is poisoned; recreate the session");
+            }
+            bail!("C++ train_multi_lora_host_i64_v2 failed");
+        }
+        Ok(MultiLoraLossReport {
+            aggregate_loss,
+            adapter_losses,
+        })
+    }
+
+    /// Get LoRA A tensor by index (for saving).
+    pub fn get_lora_a(&self, index: i64) -> Option<Tensor> {
+        let kh = get_kernels()?;
+        let ptr = unsafe { (kh.get_lora_a)(self.ptr, index) };
+        if ptr.is_null() {
+            return None;
+        }
+        Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) })
+    }
+
+    /// Get LoRA B tensor by index (for saving).
+    pub fn get_lora_b(&self, index: i64) -> Option<Tensor> {
+        let kh = get_kernels()?;
+        let ptr = unsafe { (kh.get_lora_b)(self.ptr, index) };
+        if ptr.is_null() {
+            return None;
+        }
+        Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) })
+    }
+
+    /// Inspect the native FP32 gradient accumulator for one fixed LoRA slot.
+    /// The returned tensor is a shallow handle owned by the C++ context.
+    pub fn get_lora_gradient_accumulator(&self, index: i64, is_b: bool) -> Option<Tensor> {
+        let kh = get_kernels()?;
+        let ptr =
+            unsafe { (kh.get_lora_grad_accumulator)(self.ptr, index, if is_b { 1 } else { 0 }) };
+        if ptr.is_null() {
+            return None;
+        }
+        Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) })
+    }
+
+    /// Abort an incomplete micro-batch window without changing parameters,
+    /// Adam state, or optimizer clocks. Safe to call when no window is active.
+    pub fn abort_gradient_accumulation(&self) -> Result<()> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let status = unsafe { (kh.abort_gradient_accumulation)(self.ptr) };
+        if status != 0 {
+            bail!("C++ abort_gradient_accumulation failed");
+        }
+        Ok(())
+    }
+
+    pub fn set_lora_tensor(&self, index: i64, is_b: bool, tensor: &Tensor) -> Result<()> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let status = unsafe {
+            (kh.set_lora_tensor)(
+                self.ptr,
+                index,
+                if is_b { 1 } else { 0 },
+                tensor.as_ptr() as *mut c_void,
+            )
+        };
+        if status != 0 {
+            bail!("C++ set_lora_tensor failed for slot {index}");
+        }
+        Ok(())
+    }
+
+    pub fn lora_count(&self) -> i64 {
+        self.lora_count
+    }
+
+    /// Set MTP weights on the C++ training context.
+    /// Must be called after `new()` if MTP is enabled.
+    pub fn set_mtp_weights(
+        &self,
+        weights: &std::collections::BTreeMap<String, Tensor>,
+        config: &crate::config::Qwen36RuntimeConfig,
+        expert_start: usize,
+        expert_count: usize,
+    ) -> Result<()> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+
+        let mtp_fc_ptr = get_ptr(weights, "mtp.fc.weight");
+        let mtp_pre_fc_norm_emb_ptr = get_ptr(weights, "mtp.pre_fc_norm_embedding.weight");
+        let mtp_pre_fc_norm_hidden_ptr = get_ptr(weights, "mtp.pre_fc_norm_hidden.weight");
+        let mtp_norm_ptr = get_ptr(weights, "mtp.norm.weight");
+
+        let mut mtp_weight_ptrs = build_mtp_weight_ptrs(weights, config);
+        let mtp_layer_configs = build_mtp_layer_configs(config, expert_start, expert_count);
+
+        let wp_ptr = mtp_weight_ptrs.as_ptr() as *mut *mut c_void;
+        let wp_len = mtp_weight_ptrs.len();
+        std::mem::forget(mtp_weight_ptrs);
+        let lc_ptr = mtp_layer_configs.as_ptr() as *mut c_void;
+        std::mem::forget(mtp_layer_configs);
+
+        let status = unsafe {
+            (kh.set_mtp_weights)(
+                self.ptr,
+                mtp_fc_ptr,
+                mtp_pre_fc_norm_emb_ptr,
+                mtp_pre_fc_norm_hidden_ptr,
+                mtp_norm_ptr,
+                wp_ptr,
+                wp_len as i64,
+                lc_ptr,
+                config.mtp_num_hidden_layers as i64,
+            )
+        };
+        if status != 0 {
+            bail!("C++ set_mtp_weights rejected the requested context state");
+        }
+        Ok(())
+    }
+
+    /// Enable/disable gradient checkpointing.
+    /// When enabled, layers are grouped by `group_size` and intermediate
+    /// activations are recomputed during backward instead of stored.
+    pub fn set_checkpoint(&self, enable: bool, group_size: i64) {
+        let kh = get_kernels().expect("kernels not loaded");
+        unsafe {
+            (kh.set_checkpoint)(self.ptr, if enable { 1 } else { 0 }, group_size);
+        }
+    }
+
+    /// Set NCCL communicator for Expert Parallel all-reduce.
+    /// Must be called after `new()` if EP is enabled.
+    /// comm_ptr / stream_ptr from `NcclPersistentComm::raw_comm_ptr()` / `raw_stream_ptr()`.
+    pub fn set_nccl_comm(
+        &self,
+        comm_ptr: *mut c_void,
+        stream_ptr: *mut c_void,
+        ep_rank: i32,
+        ep_world_size: i32,
+    ) {
+        let kh = get_kernels().expect("kernels not loaded");
+        unsafe {
+            (kh.set_nccl_comm)(self.ptr, comm_ptr, stream_ptr, ep_rank, ep_world_size);
+        }
+    }
+
+    /// Initialize NCCL communicator directly in C++ (preferred over set_nccl_comm).
+    /// Reads RANK/WORLD_SIZE/LOCAL_RANK from env vars.
+    /// Returns 0 on success, -1 on failure.
+    pub fn init_nccl(&self) -> i32 {
+        let kh = get_kernels().expect("kernels not loaded");
+        unsafe { (kh.init_nccl)(self.ptr) }
+    }
+
+    /// Initialize the orthogonal TP, CP, EP, DP, and PP process grid.
+    pub fn init_parallel_nccl(
+        &self,
+        rank: usize,
+        world_size: usize,
+        tp_rank: usize,
+        tp_size: usize,
+        tp_color: usize,
+        cp_rank: usize,
+        cp_size: usize,
+        cp_color: usize,
+        ep_rank: usize,
+        ep_size: usize,
+        ep_color: usize,
+        dp_rank: usize,
+        dp_size: usize,
+        dp_color: usize,
+        pp_rank: usize,
+        pp_size: usize,
+        pp_color: usize,
+    ) -> Result<()> {
+        let values = [
+            rank, world_size, tp_rank, tp_size, tp_color, cp_rank, cp_size, cp_color, ep_rank,
+            ep_size, ep_color, dp_rank, dp_size, dp_color, pp_rank, pp_size, pp_color,
+        ]
+        .map(|value| i32::try_from(value).context("parallel topology exceeds i32"));
+        let [rank, world_size, tp_rank, tp_size, tp_color, cp_rank, cp_size, cp_color, ep_rank, ep_size, ep_color, dp_rank, dp_size, dp_color, pp_rank, pp_size, pp_color] =
+            values;
+        let kh = get_kernels().expect("kernels not loaded");
+        let status = unsafe {
+            (kh.init_parallel_nccl)(
+                self.ptr,
+                rank?,
+                world_size?,
+                tp_rank?,
+                tp_size?,
+                tp_color?,
+                cp_rank?,
+                cp_size?,
+                cp_color?,
+                ep_rank?,
+                ep_size?,
+                ep_color?,
+                dp_rank?,
+                dp_size?,
+                dp_color?,
+                pp_rank?,
+                pp_size?,
+                pp_color?,
+            )
+        };
+        if status != 0 {
+            bail!("C++ parallel NCCL init failed (code {status})");
+        }
+        Ok(())
+    }
+
+    /// Attach a checkpoint shadow context to process-cached communicators.
+    /// This deliberately skips parameter broadcasts because restore replaces
+    /// every active LoRA/Adam tensor before the context can become live.
+    pub fn attach_parallel_nccl_no_sync(
+        &self,
+        rank: usize,
+        world_size: usize,
+        tp_rank: usize,
+        tp_size: usize,
+        tp_color: usize,
+        cp_rank: usize,
+        cp_size: usize,
+        cp_color: usize,
+        ep_rank: usize,
+        ep_size: usize,
+        ep_color: usize,
+        dp_rank: usize,
+        dp_size: usize,
+        dp_color: usize,
+        pp_rank: usize,
+        pp_size: usize,
+        pp_color: usize,
+    ) -> Result<()> {
+        let values = [
+            rank, world_size, tp_rank, tp_size, tp_color, cp_rank, cp_size, cp_color, ep_rank,
+            ep_size, ep_color, dp_rank, dp_size, dp_color, pp_rank, pp_size, pp_color,
+        ]
+        .map(|value| i32::try_from(value).context("parallel topology exceeds i32"));
+        let [rank, world_size, tp_rank, tp_size, tp_color, cp_rank, cp_size, cp_color, ep_rank, ep_size, ep_color, dp_rank, dp_size, dp_color, pp_rank, pp_size, pp_color] =
+            values;
+        let kh = get_kernels().expect("kernels not loaded");
+        let status = unsafe {
+            (kh.attach_parallel_nccl_no_sync)(
+                self.ptr,
+                rank?,
+                world_size?,
+                tp_rank?,
+                tp_size?,
+                tp_color?,
+                cp_rank?,
+                cp_size?,
+                cp_color?,
+                ep_rank?,
+                ep_size?,
+                ep_color?,
+                dp_rank?,
+                dp_size?,
+                dp_color?,
+                pp_rank?,
+                pp_size?,
+                pp_color?,
+            )
+        };
+        if status != 0 {
+            bail!("C++ parallel NCCL restore attach failed (code {status})");
+        }
+        Ok(())
+    }
+
+    /// Set CUDA device and force PyTorch CUDA context initialization.
+    /// Must be called before any GPU operation in worker processes.
+    pub fn set_cuda_device(device: i32) {
+        let kh = get_kernels().expect("kernels not loaded");
+        unsafe { (kh.set_cuda_device)(device) }
+    }
+
+    /// Add a new LoRA adapter. Returns adapter ID (>0) on success.
+    /// target_layers: empty = all layers
+    /// target_modules: comma-separated, e.g. "q_proj,k_proj,v_proj,o_proj". Empty = all.
+    pub fn add_lora(
+        &self,
+        rank: i64,
+        alpha: f64,
+        target_layers: &[i64],
+        target_modules: &str,
+    ) -> Result<i64> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let tl_ptr = if target_layers.is_empty() {
+            std::ptr::null()
+        } else {
+            target_layers.as_ptr()
+        };
+        let tl_len = target_layers.len() as i64;
+        let modules_c = std::ffi::CString::new(target_modules).unwrap();
+        let modules_ptr = if target_modules.is_empty() {
+            std::ptr::null()
+        } else {
+            modules_c.as_ptr()
+        };
+        let id = unsafe { (kh.add_lora_v2)(self.ptr, rank, alpha, tl_ptr, tl_len, modules_ptr) };
+        if id < 0 {
+            bail!("C++ add_lora failed");
+        }
+        Ok(id)
+    }
+
+    /// Add a dynamic adapter with a tenant-specific Adam learning rate.
+    /// The remaining hyperparameters inherit the training context defaults.
+    pub fn add_lora_with_optimizer_lr(
+        &self,
+        rank: i64,
+        alpha: f64,
+        target_layers: &[i64],
+        target_modules: &str,
+        optimizer_lr: f64,
+    ) -> Result<i64> {
+        if !optimizer_lr.is_finite() || optimizer_lr < 0.0 {
+            bail!("dynamic LoRA optimizer learning rate must be finite and non-negative");
+        }
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let add_lora = kh.add_lora_with_optimizer.ok_or_else(|| {
+            anyhow::anyhow!("loaded Qwen kernel does not support tenant optimizer overrides")
+        })?;
+        let tl_ptr = if target_layers.is_empty() {
+            std::ptr::null()
+        } else {
+            target_layers.as_ptr()
+        };
+        let modules = std::ffi::CString::new(target_modules)?;
+        let modules_ptr = if target_modules.is_empty() {
+            std::ptr::null()
+        } else {
+            modules.as_ptr()
+        };
+        let id = unsafe {
+            add_lora(
+                self.ptr,
+                rank,
+                alpha,
+                tl_ptr,
+                i64::try_from(target_layers.len()).context("target layer count exceeds i64")?,
+                modules_ptr,
+                optimizer_lr,
+            )
+        };
+        if id < 0 {
+            bail!("C++ add_lora_with_optimizer failed");
+        }
+        Ok(id)
+    }
+
+    /// Add a dynamic adapter with a complete tenant-local Adam configuration.
+    /// All selected tenants still share one transactional fused Adam launch.
+    pub fn add_lora_with_optimizer_config(
+        &self,
+        rank: i64,
+        alpha: f64,
+        target_layers: &[i64],
+        target_modules: &str,
+        optimizer: DynamicAdamConfig,
+    ) -> Result<i64> {
+        let optimizer = optimizer.validate()?;
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let add_lora = kh.add_lora_with_optimizer_v2.ok_or_else(|| {
+            anyhow::anyhow!(
+                "loaded Qwen kernel does not support complete tenant optimizer overrides"
+            )
+        })?;
+        let tl_ptr = if target_layers.is_empty() {
+            std::ptr::null()
+        } else {
+            target_layers.as_ptr()
+        };
+        let modules = std::ffi::CString::new(target_modules)?;
+        let modules_ptr = if target_modules.is_empty() {
+            std::ptr::null()
+        } else {
+            modules.as_ptr()
+        };
+        let id = unsafe {
+            add_lora(
+                self.ptr,
+                rank,
+                alpha,
+                tl_ptr,
+                i64::try_from(target_layers.len()).context("target layer count exceeds i64")?,
+                modules_ptr,
+                optimizer.lr,
+                optimizer.beta1,
+                optimizer.beta2,
+                optimizer.eps,
+            )
+        };
+        if id < 0 {
+            bail!("C++ add_lora_with_optimizer_v2 failed");
+        }
+        Ok(id)
+    }
+
+    /// Allocate a dynamic adapter for checkpoint hydration without collective
+    /// synchronization of its temporary random initialization.
+    pub fn add_lora_for_restore(
+        &self,
+        rank: i64,
+        alpha: f64,
+        target_layers: &[i64],
+        target_modules: &str,
+    ) -> Result<i64> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let tl_ptr = if target_layers.is_empty() {
+            std::ptr::null()
+        } else {
+            target_layers.as_ptr()
+        };
+        let tl_len = target_layers.len() as i64;
+        let modules_c = std::ffi::CString::new(target_modules)?;
+        let modules_ptr = if target_modules.is_empty() {
+            std::ptr::null()
+        } else {
+            modules_c.as_ptr()
+        };
+        let id = unsafe {
+            (kh.add_lora_for_restore)(self.ptr, rank, alpha, tl_ptr, tl_len, modules_ptr)
+        };
+        if id < 0 {
+            bail!("C++ restore adapter allocation failed");
+        }
+        Ok(id)
+    }
+
+    /// Allocate a dynamic adapter with a checkpointed learning rate while
+    /// suppressing synchronization of its temporary random initialization.
+    pub fn add_lora_for_restore_with_optimizer_lr(
+        &self,
+        rank: i64,
+        alpha: f64,
+        target_layers: &[i64],
+        target_modules: &str,
+        optimizer_lr: f64,
+    ) -> Result<i64> {
+        if !optimizer_lr.is_finite() || optimizer_lr < 0.0 {
+            bail!("dynamic LoRA optimizer learning rate must be finite and non-negative");
+        }
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let tl_ptr = if target_layers.is_empty() {
+            std::ptr::null()
+        } else {
+            target_layers.as_ptr()
+        };
+        let modules_c = std::ffi::CString::new(target_modules)?;
+        let modules_ptr = if target_modules.is_empty() {
+            std::ptr::null()
+        } else {
+            modules_c.as_ptr()
+        };
+        let id = unsafe {
+            (kh.add_lora_for_restore_with_optimizer)(
+                self.ptr,
+                rank,
+                alpha,
+                tl_ptr,
+                target_layers.len() as i64,
+                modules_ptr,
+                optimizer_lr,
+            )
+        };
+        if id < 0 {
+            bail!("C++ restore adapter allocation with optimizer state failed");
+        }
+        Ok(id)
+    }
+
+    /// Allocate a dynamic adapter with checkpointed tenant-local Adam
+    /// hyperparameters while suppressing temporary parameter synchronization.
+    pub fn add_lora_for_restore_with_optimizer_config(
+        &self,
+        rank: i64,
+        alpha: f64,
+        target_layers: &[i64],
+        target_modules: &str,
+        optimizer: DynamicAdamConfig,
+    ) -> Result<i64> {
+        let optimizer = optimizer.validate()?;
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let add_lora = kh.add_lora_for_restore_with_optimizer_v2.ok_or_else(|| {
+            anyhow::anyhow!("loaded Qwen kernel does not support complete tenant optimizer restore")
+        })?;
+        let tl_ptr = if target_layers.is_empty() {
+            std::ptr::null()
+        } else {
+            target_layers.as_ptr()
+        };
+        let modules = std::ffi::CString::new(target_modules)?;
+        let modules_ptr = if target_modules.is_empty() {
+            std::ptr::null()
+        } else {
+            modules.as_ptr()
+        };
+        let id = unsafe {
+            add_lora(
+                self.ptr,
+                rank,
+                alpha,
+                tl_ptr,
+                i64::try_from(target_layers.len()).context("target layer count exceeds i64")?,
+                modules_ptr,
+                optimizer.lr,
+                optimizer.beta1,
+                optimizer.beta2,
+                optimizer.eps,
+            )
+        };
+        if id < 0 {
+            bail!("C++ add_lora_for_restore_with_optimizer_v2 failed");
+        }
+        Ok(id)
+    }
+
+    /// Remove a LoRA adapter by ID.
+    pub fn remove_lora(&self, adapter_id: i64) -> Result<bool> {
+        if adapter_id == 0 {
+            return Ok(false);
+        }
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let status = unsafe { (kh.remove_lora)(self.ptr, adapter_id) };
+        if status < 0 {
+            bail!("C++ remove_lora failed for adapter {adapter_id}");
+        }
+        Ok(status != 0)
+    }
+
+    /// Restore a dynamic adapter's stable external ID from a checkpoint.
+    pub fn set_adapter_id(&self, current_id: i64, requested_id: i64) -> Result<()> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let status = unsafe { (kh.set_adapter_id)(self.ptr, current_id, requested_id) };
+        if status != 0 {
+            bail!("C++ set_adapter_id failed: {current_id} -> {requested_id}");
+        }
+        Ok(())
     }
 
     /// List all active adapter IDs.
     pub fn list_lora(&self) -> Vec<i64> {
-        let kh = match get_kernels() { Some(k) => k, None => return Vec::new() };
-        let mut ids = vec![0i64; 64];
-        let count = unsafe { (kh.list_lora)(self.ptr, ids.as_mut_ptr(), 64) };
-        ids.truncate(count as usize);
+        let mut ids = self.list_dynamic_lora();
+        if self.lora_count > 0 {
+            // ID 0 is the fixed adapter created with the training context.
+            ids.insert(0, 0);
+        }
         ids
     }
 
+    /// List only dynamic adapter IDs. Fixed adapter ID 0 is intentionally
+    /// excluded because selected multi-LoRA training accepts dynamic tenants.
+    fn list_dynamic_lora(&self) -> Vec<i64> {
+        let kh = match get_kernels() {
+            Some(k) => k,
+            None => return Vec::new(),
+        };
+        // Query the native registry size first; the old fixed 64-entry buffer
+        // silently truncated large multi-tenant registries.
+        let total = unsafe { (kh.list_lora)(self.ptr, std::ptr::null_mut(), 0) };
+        if total <= 0 {
+            return Vec::new();
+        }
+        let mut dynamic_ids = vec![0i64; total as usize];
+        let count =
+            unsafe { (kh.list_lora)(self.ptr, dynamic_ids.as_mut_ptr(), dynamic_ids.len() as i64) };
+        if count <= 0 {
+            return Vec::new();
+        }
+        dynamic_ids.truncate((count as usize).min(dynamic_ids.len()));
+        dynamic_ids
+    }
+
+    /// Returns a shallow snapshot of the adapter tensor at call time.
+    ///
+    /// Dynamic Adam may atomically replace the native registry handle. Do not
+    /// cache this tensor across a training call; fetch it again after training.
+    pub fn get_adapter_lora_tensor(
+        &self,
+        adapter_id: i64,
+        layer: i64,
+        module: &str,
+        is_b: bool,
+    ) -> Option<Tensor> {
+        let kh = get_kernels()?;
+        let module = std::ffi::CString::new(module).ok()?;
+        let ptr = unsafe {
+            (kh.get_adapter_lora_tensor)(
+                self.ptr,
+                adapter_id,
+                layer,
+                module.as_ptr(),
+                if is_b { 1 } else { 0 },
+            )
+        };
+        if ptr.is_null() {
+            return None;
+        }
+        Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) })
+    }
+
+    pub fn set_adapter_lora_tensor(
+        &self,
+        adapter_id: i64,
+        layer: i64,
+        module: &str,
+        is_b: bool,
+        tensor: &Tensor,
+    ) -> Result<()> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let module = std::ffi::CString::new(module)?;
+        let status = unsafe {
+            (kh.set_adapter_lora_tensor)(
+                self.ptr,
+                adapter_id,
+                layer,
+                module.as_ptr(),
+                if is_b { 1 } else { 0 },
+                tensor.as_ptr() as *mut c_void,
+            )
+        };
+        if status != 0 {
+            bail!(
+                "C++ set_adapter_lora_tensor failed for adapter {adapter_id}, layer {layer}, module {module:?}"
+            );
+        }
+        Ok(())
+    }
+
+    /// Returns a shallow snapshot of the optimizer tensor at call time.
+    /// Fetch it again after training because transactional Adam swaps native
+    /// registry handles instead of mutating the previously returned handle.
+    pub fn get_adapter_optimizer_tensor(
+        &self,
+        adapter_id: i64,
+        layer: i64,
+        module: &str,
+        is_b: bool,
+        is_v: bool,
+    ) -> Option<Tensor> {
+        let kh = get_kernels()?;
+        let module = std::ffi::CString::new(module).ok()?;
+        let ptr = unsafe {
+            (kh.get_adapter_optimizer_tensor)(
+                self.ptr,
+                adapter_id,
+                layer,
+                module.as_ptr(),
+                if is_b { 1 } else { 0 },
+                if is_v { 1 } else { 0 },
+            )
+        };
+        if ptr.is_null() {
+            return None;
+        }
+        Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) })
+    }
+
+    /// Export one dynamic Adam tensor as an owned CPU snapshot for checkpointing.
+    /// `None` is valid only for a tenant that has not completed an optimizer step.
+    pub fn export_adapter_optimizer_tensor_cpu(
+        &self,
+        adapter_id: i64,
+        layer: i64,
+        module: &str,
+        is_b: bool,
+        is_v: bool,
+    ) -> Result<Option<Tensor>> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let module = std::ffi::CString::new(module)?;
+        let mut ptr = std::ptr::null_mut();
+        let status = unsafe {
+            (kh.export_adapter_optimizer_tensor_cpu)(
+                self.ptr,
+                adapter_id,
+                layer,
+                module.as_ptr(),
+                if is_b { 1 } else { 0 },
+                if is_v { 1 } else { 0 },
+                &mut ptr,
+            )
+        };
+        match status {
+            0 if !ptr.is_null() => {
+                let tensor = unsafe { Tensor::clone_from_ptr(ptr as *mut _) };
+                unsafe { (kh.free_tensor)(ptr) };
+                Ok(Some(tensor))
+            }
+            1 if ptr.is_null() => Ok(None),
+            _ => {
+                if !ptr.is_null() {
+                    unsafe { (kh.free_tensor)(ptr) };
+                }
+                bail!(
+                    "C++ dynamic optimizer CPU export failed for adapter {adapter_id}, layer {layer}, module {module:?}, status={status}"
+                );
+            }
+        }
+    }
+
+    /// Atomically install a complete stage-local Adam checkpoint on CPU.
+    pub fn import_adapter_optimizer_state_host(
+        &self,
+        adapter_id: i64,
+        optimizer_step: i64,
+        global_layers: &[i64],
+        modules: &[&str],
+        adam_m: &[Tensor],
+        adam_v: &[Tensor],
+    ) -> Result<()> {
+        if optimizer_step < 0 {
+            bail!("dynamic Adam host import step must be non-negative");
+        }
+        if global_layers.len() != modules.len() ||
+            adam_m.len() != global_layers.len().saturating_mul(2) ||
+            adam_v.len() != global_layers.len().saturating_mul(2)
+        {
+            bail!(
+                "dynamic Adam host import layout mismatch: layers={} modules={} m={} v={}",
+                global_layers.len(), modules.len(), adam_m.len(), adam_v.len()
+            );
+        }
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let modules = modules
+            .iter()
+            .map(|module| std::ffi::CString::new(*module))
+            .collect::<std::result::Result<Vec<_>, _>>()?;
+        let module_ptrs = modules
+            .iter()
+            .map(|module| module.as_ptr())
+            .collect::<Vec<_>>();
+        let mut state_ptrs = Vec::with_capacity(global_layers.len().saturating_mul(4));
+        for slot in 0..global_layers.len() {
+            state_ptrs.push(adam_m[slot * 2].as_ptr() as *mut c_void);
+            state_ptrs.push(adam_v[slot * 2].as_ptr() as *mut c_void);
+            state_ptrs.push(adam_m[slot * 2 + 1].as_ptr() as *mut c_void);
+            state_ptrs.push(adam_v[slot * 2 + 1].as_ptr() as *mut c_void);
+        }
+        let status = unsafe {
+            (kh.import_adapter_optimizer_state_host)(
+                self.ptr,
+                adapter_id,
+                global_layers.as_ptr(),
+                module_ptrs.as_ptr(),
+                state_ptrs.as_ptr(),
+                i64::try_from(global_layers.len()).context("optimizer slot count exceeds i64")?,
+                optimizer_step,
+            )
+        };
+        if status != 0 {
+            bail!("C++ dynamic Adam host import failed for adapter {adapter_id}");
+        }
+        Ok(())
+    }
+
+    pub fn get_adapter_optimizer_resident_count(&self, adapter_id: i64) -> Result<i64> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let count = unsafe { (kh.get_adapter_optimizer_resident_count)(self.ptr, adapter_id) };
+        if count < 0 {
+            bail!("C++ dynamic Adam resident count failed for adapter {adapter_id}");
+        }
+        Ok(count)
+    }
+
+    pub fn set_adapter_optimizer_tensor(
+        &self,
+        adapter_id: i64,
+        layer: i64,
+        module: &str,
+        is_b: bool,
+        is_v: bool,
+        tensor: &Tensor,
+    ) -> Result<()> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let module = std::ffi::CString::new(module)?;
+        let status = unsafe {
+            (kh.set_adapter_optimizer_tensor)(
+                self.ptr,
+                adapter_id,
+                layer,
+                module.as_ptr(),
+                if is_b { 1 } else { 0 },
+                if is_v { 1 } else { 0 },
+                tensor.as_ptr() as *mut c_void,
+            )
+        };
+        if status != 0 {
+            bail!(
+                "C++ set_adapter_optimizer_tensor failed for adapter {adapter_id}, layer {layer}, module {module:?}"
+            );
+        }
+        Ok(())
+    }
+
+    /// Return the independent Adam bias-correction clock for a dynamic tenant.
+    pub fn get_adapter_step_count(&self, adapter_id: i64) -> Result<i64> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let step = unsafe { (kh.get_adapter_step_count)(self.ptr, adapter_id) };
+        if step < 0 {
+            bail!("C++ get_adapter_step_count failed for adapter {adapter_id}");
+        }
+        Ok(step)
+    }
+
+    /// A failed native dynamic-LoRA recovery quarantines the whole context.
+    /// ABI28 libraries predating the health symbol are treated as healthy.
+    pub fn is_healthy(&self) -> bool {
+        let Some(kh) = get_kernels() else {
+            return false;
+        };
+        kh.get_context_health
+            .map(|health| unsafe { health(self.ptr) == 0 })
+            .unwrap_or(true)
+    }
+
+    pub fn validate_adapter_steps(
+        &self,
+        adapter_ids: &[i64],
+        expected_steps: &[u64],
+    ) -> Result<()> {
+        if adapter_ids.is_empty() || adapter_ids.len() != expected_steps.len() {
+            bail!(
+                "expected_steps length {} must match non-empty adapter_ids length {}",
+                expected_steps.len(),
+                adapter_ids.len()
+            );
+        }
+        let expected_steps = expected_steps
+            .iter()
+            .map(|step| i64::try_from(*step).context("expected adapter step exceeds i64"))
+            .collect::<Result<Vec<_>>>()?;
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let validate = kh.validate_adapter_steps.ok_or_else(|| {
+            anyhow::anyhow!("native library does not export adapter step validation")
+        })?;
+        let status = unsafe {
+            validate(
+                self.ptr,
+                adapter_ids.as_ptr(),
+                expected_steps.as_ptr(),
+                i32::try_from(adapter_ids.len()).context("adapter count exceeds i32")?,
+            )
+        };
+        if status != 0 {
+            bail!("C++ dynamic adapter step validation failed");
+        }
+        Ok(())
+    }
+
+    /// Restore a dynamic tenant's Adam bias-correction clock.
+    pub fn set_adapter_step_count(&self, adapter_id: i64, step_count: i64) -> Result<()> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let status = unsafe { (kh.set_adapter_step_count)(self.ptr, adapter_id, step_count) };
+        if status != 0 {
+            bail!("C++ set_adapter_step_count failed for adapter {adapter_id}");
+        }
+        Ok(())
+    }
+
     /// Eval step: forward + loss, no backward, no Adam update.
-    pub fn eval_step(&self, input_ids: &Tensor, target_mask: &Tensor, attention_mask: &Tensor) -> Result<f64> {
+    pub fn eval_step(
+        &self,
+        input_ids: &Tensor,
+        target_mask: &Tensor,
+        attention_mask: &Tensor,
+    ) -> Result<f64> {
         let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
         let loss = unsafe {
             (kh.eval_step)(
@@ -572,32 +2727,111 @@ impl CppTrainingContext {
                 attention_mask.as_ptr() as *mut _,
             )
         };
-        if loss < 0.0 {
+        if !loss.is_finite() || loss < 0.0 {
             bail!("C++ eval_step failed");
         }
         Ok(loss)
     }
 
+    /// Evaluate a borrowed host int64 batch without constructing tensors in Rust.
+    pub fn eval_step_host_i64(
+        &self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+    ) -> Result<f64> {
+        validate_host_batch(input_ids, target_mask, attention_mask, batch_size, seq_len)?;
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let loss = unsafe {
+            (kh.eval_step_host_i64)(
+                self.ptr,
+                input_ids.as_ptr(),
+                target_mask.as_ptr(),
+                attention_mask.as_ptr(),
+                i64::try_from(batch_size).context("host batch_size exceeds i64")?,
+                i64::try_from(seq_len).context("host seq_len exceeds i64")?,
+            )
+        };
+        if !loss.is_finite() || loss < 0.0 {
+            bail!("C++ eval_step_host_i64 failed");
+        }
+        Ok(loss)
+    }
+
+    /// Evaluate selected dynamic tenants independently and return one loss per
+    /// requested adapter. The native implementation scopes one registry entry
+    /// at a time, so heterogeneous tenant ranks remain isolated.
+    pub fn eval_multi_lora_host_i64(
+        &self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+        adapter_ids: &[i64],
+    ) -> Result<Vec<f64>> {
+        validate_host_batch(input_ids, target_mask, attention_mask, batch_size, seq_len)?;
+        if adapter_ids.is_empty() {
+            bail!("selected multi-LoRA eval requires at least one adapter ID");
+        }
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let eval = kh.eval_multi_lora_host_i64.ok_or_else(|| {
+            anyhow::anyhow!("native library does not export selected multi-LoRA eval")
+        })?;
+        let mut losses = vec![f64::NAN; adapter_ids.len()];
+        let status = unsafe {
+            eval(
+                self.ptr,
+                input_ids.as_ptr(),
+                target_mask.as_ptr(),
+                attention_mask.as_ptr(),
+                i64::try_from(batch_size).context("host batch_size exceeds i64")?,
+                i64::try_from(seq_len).context("host seq_len exceeds i64")?,
+                adapter_ids.as_ptr(),
+                i32::try_from(adapter_ids.len()).context("adapter count exceeds i32")?,
+                losses.as_mut_ptr(),
+                i32::try_from(losses.len()).context("adapter loss capacity exceeds i32")?,
+            )
+        };
+        if status != 0 || losses.iter().any(|loss| !loss.is_finite() || *loss < 0.0) {
+            if !self.is_healthy() {
+                bail!("native dynamic LoRA context is poisoned; recreate the session");
+            }
+            bail!("C++ selected multi-LoRA eval failed");
+        }
+        Ok(losses)
+    }
+
     /// Get current training step count.
     pub fn get_step_count(&self) -> i64 {
-        let kh = match get_kernels() { Some(k) => k, None => return 0 };
+        let kh = match get_kernels() {
+            Some(k) => k,
+            None => return 0,
+        };
         unsafe { (kh.get_step_count)(self.ptr) }
     }
 
+    /// Restore the native Adam bias-correction step from a checkpoint.
+    pub fn set_step_count(&self, step_count: i64) -> Result<()> {
+        let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
+        let status = unsafe { (kh.set_step_count)(self.ptr, step_count) };
+        if status != 0 {
+            bail!("C++ set_step_count failed for step {step_count}");
+        }
+        Ok(())
+    }
+
     /// Export Adam optimizer state (m and v vectors).
     /// Returns (m_tensors, v_tensors) — owned copies on CPU.
     pub fn export_optimizer_state(&self) -> Result<(Vec<Tensor>, Vec<Tensor>)> {
         let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
-        let count = self.lora_count * 2;  // m and v per LoRA param (a+b)
+        let count = self.lora_count * 2; // m and v per LoRA param (a+b)
         let mut m_ptrs: Vec<*mut c_void> = vec![std::ptr::null_mut(); count as usize];
         let mut v_ptrs: Vec<*mut c_void> = vec![std::ptr::null_mut(); count as usize];
         let actual = unsafe {
-            (kh.export_optimizer)(
-                self.ptr,
-                m_ptrs.as_mut_ptr(),
-                v_ptrs.as_mut_ptr(),
-                count,
-            )
+            (kh.export_optimizer)(self.ptr, m_ptrs.as_mut_ptr(), v_ptrs.as_mut_ptr(), count)
         };
         let mut m_tensors = Vec::new();
         let mut v_tensors = Vec::new();
@@ -613,11 +2847,21 @@ impl CppTrainingContext {
     }
 
     /// Import Adam optimizer state (m and v vectors).
-    pub fn import_optimizer_state(&self, m_tensors: &[Tensor], v_tensors: &[Tensor]) -> Result<i64> {
+    pub fn import_optimizer_state(
+        &self,
+        m_tensors: &[Tensor],
+        v_tensors: &[Tensor],
+    ) -> Result<i64> {
         let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?;
         let count = m_tensors.len().min(v_tensors.len());
-        let m_ptrs: Vec<*mut c_void> = m_tensors.iter().map(|t| t.as_ptr() as *mut c_void).collect();
-        let v_ptrs: Vec<*mut c_void> = v_tensors.iter().map(|t| t.as_ptr() as *mut c_void).collect();
+        let m_ptrs: Vec<*mut c_void> = m_tensors
+            .iter()
+            .map(|t| t.as_ptr() as *mut c_void)
+            .collect();
+        let v_ptrs: Vec<*mut c_void> = v_tensors
+            .iter()
+            .map(|t| t.as_ptr() as *mut c_void)
+            .collect();
         let imported = unsafe {
             (kh.import_optimizer)(
                 self.ptr,
@@ -626,10 +2870,40 @@ impl CppTrainingContext {
                 count as i64,
             )
         };
+        if imported < 0 {
+            anyhow::bail!("native Adam optimizer state import failed");
+        }
         Ok(imported)
     }
 }
 
+fn validate_host_batch(
+    input_ids: &[i64],
+    target_mask: &[i64],
+    attention_mask: &[i64],
+    batch_size: usize,
+    seq_len: usize,
+) -> Result<()> {
+    if batch_size == 0 || seq_len == 0 {
+        bail!("host batch dimensions must be positive");
+    }
+    let expected = batch_size
+        .checked_mul(seq_len)
+        .context("host batch shape overflows usize")?;
+    for (name, actual) in [
+        ("input_ids", input_ids.len()),
+        ("target_mask", target_mask.len()),
+        ("attention_mask", attention_mask.len()),
+    ] {
+        if actual != expected {
+            bail!(
+                "host {name} length {actual} does not match batch_size={batch_size} * seq_len={seq_len}"
+            );
+        }
+    }
+    Ok(())
+}
+
 impl Drop for CppTrainingContext {
     fn drop(&mut self) {
         if let Some(kh) = get_kernels() {
@@ -639,3 +2913,300 @@ impl Drop for CppTrainingContext {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{
+        shard_dense_mlp_weight_for_tp, shard_full_attention_weight_for_tp,
+        shard_linear_attention_weight_for_tp, shard_moe_mlp_weight_for_tp,
+        shard_vocab_weight_for_tp,
+    };
+    use tch::{Kind, Tensor};
+
+    #[test]
+    fn dense_mlp_tp_shards_matching_intermediate_axes() {
+        let gate = Tensor::arange(48, (Kind::Float, tch::Device::Cpu)).reshape([12, 4]);
+        let down = Tensor::arange(48, (Kind::Float, tch::Device::Cpu)).reshape([4, 12]);
+        let gate_rank_one =
+            shard_dense_mlp_weight_for_tp("model.layers.0.mlp.gate_proj.weight", &gate, 2, 1)
+                .unwrap()
+                .unwrap();
+        let down_rank_one =
+            shard_dense_mlp_weight_for_tp("model.layers.0.mlp.down_proj.weight", &down, 2, 1)
+                .unwrap()
+                .unwrap();
+        assert_eq!(gate_rank_one.size(), [6, 4]);
+        assert_eq!(down_rank_one.size(), [4, 6]);
+        assert_eq!(gate_rank_one.double_value(&[0, 0]), 24.0);
+        assert_eq!(down_rank_one.double_value(&[0, 0]), 6.0);
+    }
+
+    #[test]
+    fn vocabulary_tp_shards_embedding_and_head_rows() {
+        let tensor = Tensor::arange(48, (Kind::Float, tch::Device::Cpu)).reshape([12, 4]);
+        let embed = shard_vocab_weight_for_tp(
+            "model.language_model.embed_tokens.weight",
+            &tensor,
+            12,
+            2,
+            1,
+        )
+        .unwrap()
+        .unwrap();
+        let head = shard_vocab_weight_for_tp("lm_head.weight", &tensor, 12, 2, 1)
+            .unwrap()
+            .unwrap();
+        assert_eq!(embed.size(), [6, 4]);
+        assert_eq!(head.size(), [6, 4]);
+        assert_eq!(embed.double_value(&[0, 0]), 24.0);
+        assert_eq!(head.double_value(&[0, 0]), 24.0);
+    }
+
+    #[test]
+    fn vocabulary_tp_ignores_other_weights_and_rejects_invalid_layouts() {
+        let tensor = Tensor::zeros([12, 4], (Kind::Float, tch::Device::Cpu));
+        assert!(
+            shard_vocab_weight_for_tp("model.norm.weight", &tensor, 12, 2, 0)
+                .unwrap()
+                .is_none()
+        );
+        assert!(shard_vocab_weight_for_tp("lm_head.weight", &tensor, 11, 2, 0).is_err());
+        assert!(shard_vocab_weight_for_tp(
+            "model.embed_tokens.weight",
+            &Tensor::zeros([11, 4], (Kind::Float, tch::Device::Cpu)),
+            12,
+            2,
+            0,
+        )
+        .is_err());
+    }
+
+    #[test]
+    fn dense_mlp_tp_ignores_non_mlp_weights_and_rejects_bad_shapes() {
+        let tensor = Tensor::zeros([5, 4], (Kind::Float, tch::Device::Cpu));
+        assert!(shard_dense_mlp_weight_for_tp(
+            "model.layers.0.self_attn.q_proj.weight",
+            &tensor,
+            2,
+            0
+        )
+        .unwrap()
+        .is_none());
+        assert!(
+            shard_dense_mlp_weight_for_tp("model.layers.0.mlp.up_proj.weight", &tensor, 2, 0)
+                .is_err()
+        );
+    }
+
+    #[test]
+    fn moe_tp_repacks_each_gate_up_half_before_concatenation() {
+        let packed = Tensor::arange(96, (Kind::Float, tch::Device::Cpu)).reshape([2, 12, 4]);
+        let rank_zero =
+            shard_moe_mlp_weight_for_tp("model.layers.0.mlp.experts.gate_up_proj", &packed, 2, 0)
+                .unwrap()
+                .unwrap();
+        let rank_one =
+            shard_moe_mlp_weight_for_tp("model.layers.0.mlp.experts.gate_up_proj", &packed, 2, 1)
+                .unwrap()
+                .unwrap();
+        assert_eq!(rank_zero.size(), [2, 6, 4]);
+        assert_eq!(rank_one.size(), [2, 6, 4]);
+        assert_eq!(rank_zero.double_value(&[0, 3, 0]), 24.0);
+        assert_eq!(rank_one.double_value(&[0, 0, 0]), 12.0);
+        assert_eq!(rank_one.double_value(&[0, 3, 0]), 36.0);
+
+        let rebuilt_gate = Tensor::cat(&[&rank_zero.narrow(1, 0, 3), &rank_one.narrow(1, 0, 3)], 1);
+        let rebuilt_up = Tensor::cat(&[&rank_zero.narrow(1, 3, 3), &rank_one.narrow(1, 3, 3)], 1);
+        let rebuilt = Tensor::cat(&[&rebuilt_gate, &rebuilt_up], 1);
+        assert_eq!(
+            rebuilt
+                .f_sub(&packed)
+                .unwrap()
+                .abs()
+                .max()
+                .double_value(&[]),
+            0.0
+        );
+    }
+
+    #[test]
+    fn moe_tp_shards_routed_and_shared_down_input_axes() {
+        let routed = Tensor::arange(96, (Kind::Float, tch::Device::Cpu)).reshape([2, 4, 12]);
+        let shared = Tensor::arange(48, (Kind::Float, tch::Device::Cpu)).reshape([4, 12]);
+        let routed_rank_one =
+            shard_moe_mlp_weight_for_tp("model.layers.0.mlp.experts.down_proj", &routed, 2, 1)
+                .unwrap()
+                .unwrap();
+        let shared_rank_one = shard_moe_mlp_weight_for_tp(
+            "model.layers.0.mlp.shared_expert.down_proj.weight",
+            &shared,
+            2,
+            1,
+        )
+        .unwrap()
+        .unwrap();
+        assert_eq!(routed_rank_one.size(), [2, 4, 6]);
+        assert_eq!(shared_rank_one.size(), [4, 6]);
+        assert_eq!(routed_rank_one.double_value(&[0, 0, 0]), 6.0);
+        assert_eq!(shared_rank_one.double_value(&[0, 0]), 6.0);
+    }
+
+    #[test]
+    fn moe_tp_rejects_invalid_packed_and_intermediate_shapes() {
+        let odd_packed = Tensor::zeros([2, 10, 4], (Kind::Float, tch::Device::Cpu));
+        assert!(shard_moe_mlp_weight_for_tp(
+            "model.layers.0.mlp.experts.gate_up_proj",
+            &odd_packed,
+            2,
+            0,
+        )
+        .is_err());
+        let down = Tensor::zeros([2, 4, 5], (Kind::Float, tch::Device::Cpu));
+        assert!(
+            shard_moe_mlp_weight_for_tp("model.layers.0.mlp.experts.down_proj", &down, 2, 0,)
+                .is_err()
+        );
+    }
+
+    #[test]
+    fn full_attention_tp_shards_head_and_output_axes() {
+        let q = Tensor::arange(96, (Kind::Float, tch::Device::Cpu)).reshape([24, 4]);
+        let o = Tensor::arange(48, (Kind::Float, tch::Device::Cpu)).reshape([4, 12]);
+        let q_rank_one =
+            shard_full_attention_weight_for_tp("model.layers.0.self_attn.q_proj.weight", &q, 2, 1)
+                .unwrap()
+                .unwrap();
+        let o_rank_one =
+            shard_full_attention_weight_for_tp("model.layers.0.self_attn.o_proj.weight", &o, 2, 1)
+                .unwrap()
+                .unwrap();
+        assert_eq!(q_rank_one.size(), [12, 4]);
+        assert_eq!(o_rank_one.size(), [4, 6]);
+        assert_eq!(q_rank_one.double_value(&[0, 0]), 48.0);
+        assert_eq!(o_rank_one.double_value(&[0, 0]), 6.0);
+    }
+
+    #[test]
+    fn full_attention_tp_ignores_other_weights_and_rejects_bad_shapes() {
+        let tensor = Tensor::zeros([5, 4], (Kind::Float, tch::Device::Cpu));
+        assert!(shard_full_attention_weight_for_tp(
+            "model.layers.0.mlp.gate_proj.weight",
+            &tensor,
+            2,
+            0,
+        )
+        .unwrap()
+        .is_none());
+        assert!(shard_full_attention_weight_for_tp(
+            "model.layers.0.self_attn.q_proj.weight",
+            &tensor,
+            2,
+            0,
+        )
+        .is_err());
+    }
+
+    #[test]
+    fn linear_attention_tp_preserves_flat_qkv_and_conv_segments() {
+        let qkv = Tensor::arange(48, (Kind::Float, tch::Device::Cpu)).reshape([12, 4]);
+        let conv = Tensor::arange(24, (Kind::Float, tch::Device::Cpu)).reshape([12, 1, 2]);
+        let qkv_rank_one = shard_linear_attention_weight_for_tp(
+            "model.layers.0.linear_attn.in_proj_qkv.weight",
+            &qkv,
+            2,
+            1,
+            2,
+            2,
+            4,
+            1,
+        )
+        .unwrap()
+        .unwrap();
+        let conv_rank_one = shard_linear_attention_weight_for_tp(
+            "model.layers.0.linear_attn.conv1d.weight",
+            &conv,
+            2,
+            1,
+            2,
+            2,
+            4,
+            1,
+        )
+        .unwrap()
+        .unwrap();
+        assert_eq!(qkv_rank_one.size(), [6, 4]);
+        assert_eq!(conv_rank_one.size(), [6, 1, 2]);
+        // Global rows are Q=[0..4], K=[4..8], V=[8..12]. Rank 1 owns the
+        // second half of each segment and repacks them as local [Q|K|V].
+        assert_eq!(qkv_rank_one.double_value(&[0, 0]), 8.0);
+        assert_eq!(qkv_rank_one.double_value(&[2, 0]), 24.0);
+        assert_eq!(qkv_rank_one.double_value(&[4, 0]), 40.0);
+        assert_eq!(conv_rank_one.double_value(&[0, 0, 0]), 4.0);
+        assert_eq!(conv_rank_one.double_value(&[2, 0, 0]), 12.0);
+        assert_eq!(conv_rank_one.double_value(&[4, 0, 0]), 20.0);
+    }
+
+    #[test]
+    fn linear_attention_tp_shards_value_head_tensors_and_output_columns() {
+        let z = Tensor::arange(32, (Kind::Float, tch::Device::Cpu)).reshape([8, 4]);
+        let a = Tensor::arange(16, (Kind::Float, tch::Device::Cpu)).reshape([4, 4]);
+        let a_log = Tensor::arange(4, (Kind::Float, tch::Device::Cpu));
+        let out = Tensor::arange(32, (Kind::Float, tch::Device::Cpu)).reshape([4, 8]);
+        let shard = |name: &str, tensor: &Tensor| {
+            shard_linear_attention_weight_for_tp(name, tensor, 2, 1, 2, 2, 4, 2)
+                .unwrap()
+                .unwrap()
+        };
+        let z = shard("model.layers.0.linear_attn.in_proj_z.weight", &z);
+        let a = shard("model.layers.0.linear_attn.in_proj_a.weight", &a);
+        let a_log = shard("model.layers.0.linear_attn.A_log", &a_log);
+        let out = shard("model.layers.0.linear_attn.out_proj.weight", &out);
+        assert_eq!(z.size(), [4, 4]);
+        assert_eq!(a.size(), [2, 4]);
+        assert_eq!(a_log.size(), [2]);
+        assert_eq!(out.size(), [4, 4]);
+        assert_eq!(z.double_value(&[0, 0]), 16.0);
+        assert_eq!(a.double_value(&[0, 0]), 8.0);
+        assert_eq!(a_log.double_value(&[0]), 2.0);
+        assert_eq!(out.double_value(&[0, 0]), 4.0);
+    }
+
+    #[test]
+    fn linear_attention_tp_rejects_invalid_groups_and_shapes() {
+        let tensor = Tensor::zeros([12, 4], (Kind::Float, tch::Device::Cpu));
+        assert!(shard_linear_attention_weight_for_tp(
+            "model.layers.0.linear_attn.norm.weight",
+            &tensor,
+            2,
+            0,
+            2,
+            2,
+            4,
+            1,
+        )
+        .unwrap()
+        .is_none());
+        assert!(shard_linear_attention_weight_for_tp(
+            "model.layers.0.linear_attn.in_proj_qkv.weight",
+            &tensor,
+            2,
+            0,
+            3,
+            2,
+            4,
+            1,
+        )
+        .is_err());
+        assert!(shard_linear_attention_weight_for_tp(
+            "model.layers.0.linear_attn.in_proj_qkv.weight",
+            &Tensor::zeros([11, 4], (Kind::Float, tch::Device::Cpu)),
+            2,
+            0,
+            2,
+            2,
+            4,
+            1,
+        )
+        .is_err());
+    }
+}
diff --git a/crates/rustrain-qwen3-6/src/lib.rs b/crates/rustrain-qwen3-6/src/lib.rs
index be165274..56d8edc3 100644
--- a/crates/rustrain-qwen3-6/src/lib.rs
+++ b/crates/rustrain-qwen3-6/src/lib.rs
@@ -1,11 +1,13 @@
 #![allow(unused_imports)]
 #![allow(dead_code)]
 
+pub mod checkpoint;
 pub mod config;
 pub mod kernel;
 pub mod lora;
 pub mod model;
 pub mod mtp;
+pub mod pipeline;
 pub mod session;
 pub mod sft;
 pub mod vision;
diff --git a/crates/rustrain-qwen3-6/src/lora.rs b/crates/rustrain-qwen3-6/src/lora.rs
index f88c13ee..ce4a0c11 100644
--- a/crates/rustrain-qwen3-6/src/lora.rs
+++ b/crates/rustrain-qwen3-6/src/lora.rs
@@ -1,14 +1,18 @@
 //! Qwen3.6 LoRA adapter registry — VarStore-backed, stores tensors directly.
 
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
+use std::fs;
+use std::path::{Path, PathBuf};
 
-use anyhow::{Context, Result, bail};
+use anyhow::{bail, Context, Result};
+use serde::{Deserialize, Serialize};
 use tch::{nn, Kind, Tensor};
 use tracing::info;
 
-use crate::config::Qwen36RuntimeConfig;
+use crate::config::{LayerType, Qwen36RuntimeConfig};
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
 pub enum Qwen36LoraTargetModule {
     QProj,
     KProj,
@@ -16,10 +20,95 @@ pub enum Qwen36LoraTargetModule {
     OProj,
     InProjQkv,
     InProjZ,
+    InProjA,
+    InProjB,
     OutProj,
+    GateProj,
+    UpProj,
+    DownProj,
     SharedGateProj,
     SharedUpProj,
     SharedDownProj,
+    ExpertsGateUpProj,
+    ExpertsDownProj,
+}
+
+#[derive(Debug, Clone)]
+pub struct Qwen36NativeLoraSlot {
+    pub index: usize,
+    pub layer: usize,
+    pub module: Qwen36LoraTargetModule,
+    pub active: bool,
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct Qwen36AdapterConfig {
+    pub format_version: u32,
+    pub peft_type: String,
+    pub task_type: String,
+    pub base_model_name_or_path: String,
+    pub rustrain_architecture: String,
+    pub model_family: String,
+    pub r: i64,
+    pub lora_alpha: f64,
+    pub target_layers: Vec<usize>,
+    pub target_modules: Vec<String>,
+    pub adapter_dtype: String,
+    pub bias: String,
+    pub inference_mode: bool,
+}
+
+pub struct Qwen36AdapterArtifact {
+    pub config: Qwen36AdapterConfig,
+    pub tensors: BTreeMap<String, Tensor>,
+}
+
+#[derive(Debug, Serialize)]
+struct PeftAdapterConfig<'a> {
+    peft_type: &'static str,
+    task_type: &'static str,
+    base_model_name_or_path: &'a str,
+    r: i64,
+    lora_alpha: f64,
+    lora_dropout: f64,
+    target_modules: &'a [String],
+    layers_to_transform: &'a [usize],
+    layers_pattern: &'static str,
+    bias: &'static str,
+    fan_in_fan_out: bool,
+    inference_mode: bool,
+}
+
+#[derive(Debug, Deserialize)]
+struct PeftAdapterConfigOwned {
+    #[serde(default = "default_lora_type")]
+    peft_type: String,
+    #[serde(default = "default_task_type")]
+    task_type: String,
+    #[serde(default)]
+    base_model_name_or_path: String,
+    r: i64,
+    lora_alpha: f64,
+    #[serde(default)]
+    target_modules: Vec<String>,
+    #[serde(default)]
+    layers_to_transform: Vec<usize>,
+    #[serde(default = "default_bias")]
+    bias: String,
+    #[serde(default)]
+    inference_mode: bool,
+}
+
+fn default_lora_type() -> String {
+    "LORA".to_string()
+}
+
+fn default_task_type() -> String {
+    "CAUSAL_LM".to_string()
+}
+
+fn default_bias() -> String {
+    "none".to_string()
 }
 
 impl Qwen36LoraTargetModule {
@@ -31,10 +120,17 @@ impl Qwen36LoraTargetModule {
             "o_proj" => Ok(Self::OProj),
             "in_proj_qkv" => Ok(Self::InProjQkv),
             "in_proj_z" => Ok(Self::InProjZ),
+            "in_proj_a" => Ok(Self::InProjA),
+            "in_proj_b" => Ok(Self::InProjB),
             "out_proj" => Ok(Self::OutProj),
+            "gate_proj" => Ok(Self::GateProj),
+            "up_proj" => Ok(Self::UpProj),
+            "down_proj" => Ok(Self::DownProj),
             "shared_gate_proj" => Ok(Self::SharedGateProj),
             "shared_up_proj" => Ok(Self::SharedUpProj),
             "shared_down_proj" => Ok(Self::SharedDownProj),
+            "experts_gate_up_proj" => Ok(Self::ExpertsGateUpProj),
+            "experts_down_proj" => Ok(Self::ExpertsDownProj),
             other => bail!("unknown LoRA target module: {other}"),
         }
     }
@@ -47,14 +143,536 @@ impl Qwen36LoraTargetModule {
             Self::OProj => "self_attn.o_proj",
             Self::InProjQkv => "linear_attn.in_proj_qkv",
             Self::InProjZ => "linear_attn.in_proj_z",
+            Self::InProjA => "linear_attn.in_proj_a",
+            Self::InProjB => "linear_attn.in_proj_b",
             Self::OutProj => "linear_attn.out_proj",
+            Self::GateProj => "mlp.gate_proj",
+            Self::UpProj => "mlp.up_proj",
+            Self::DownProj => "mlp.down_proj",
             Self::SharedGateProj => "mlp.shared_expert.gate_proj",
             Self::SharedUpProj => "mlp.shared_expert.up_proj",
             Self::SharedDownProj => "mlp.shared_expert.down_proj",
+            Self::ExpertsGateUpProj => "mlp.experts.gate_up_proj",
+            Self::ExpertsDownProj => "mlp.experts.down_proj",
+        }
+    }
+
+    /// Stable C++ projection identifier used by the native training context.
+    pub fn cpp_name(&self) -> &'static str {
+        match self {
+            Self::QProj => "q_proj",
+            Self::KProj => "k_proj",
+            Self::VProj => "v_proj",
+            Self::OProj => "o_proj",
+            Self::InProjQkv => "in_proj_qkv",
+            Self::InProjZ => "in_proj_z",
+            Self::InProjA => "in_proj_a",
+            Self::InProjB => "in_proj_b",
+            Self::OutProj => "out_proj",
+            Self::GateProj => "gate_proj",
+            Self::UpProj => "up_proj",
+            Self::DownProj => "down_proj",
+            Self::SharedGateProj => "shared_gate_proj",
+            Self::SharedUpProj => "shared_up_proj",
+            Self::SharedDownProj => "shared_down_proj",
+            Self::ExpertsGateUpProj => "experts_gate_up_proj",
+            Self::ExpertsDownProj => "experts_down_proj",
+        }
+    }
+}
+
+pub fn native_lora_slots(
+    config: &Qwen36RuntimeConfig,
+    lora_config: &Qwen36LoraConfig,
+) -> Vec<Qwen36NativeLoraSlot> {
+    let target_layers = lora_config
+        .target_layers
+        .iter()
+        .copied()
+        .collect::<BTreeSet<_>>();
+    let target_modules = lora_config
+        .target_modules
+        .iter()
+        .copied()
+        .collect::<BTreeSet<_>>();
+    let all_layers = target_layers.is_empty();
+    let all_modules = target_modules.is_empty();
+    let mut slots = Vec::new();
+
+    for (layer, layer_type) in config.layer_types.iter().enumerate() {
+        let modules: &[Qwen36LoraTargetModule] = match layer_type {
+            LayerType::FullAttention => &[
+                Qwen36LoraTargetModule::QProj,
+                Qwen36LoraTargetModule::KProj,
+                Qwen36LoraTargetModule::VProj,
+                Qwen36LoraTargetModule::OProj,
+            ],
+            LayerType::LinearAttention => &[
+                Qwen36LoraTargetModule::InProjQkv,
+                Qwen36LoraTargetModule::InProjZ,
+                Qwen36LoraTargetModule::InProjA,
+                Qwen36LoraTargetModule::InProjB,
+                Qwen36LoraTargetModule::OutProj,
+            ],
+        };
+        for &module in modules {
+            slots.push(Qwen36NativeLoraSlot {
+                index: slots.len(),
+                layer,
+                module,
+                active: (all_layers || target_layers.contains(&layer))
+                    && (all_modules || target_modules.contains(&module)),
+            });
+        }
+        let mlp_modules: &[Qwen36LoraTargetModule] = if config.is_moe {
+            &[
+                Qwen36LoraTargetModule::SharedGateProj,
+                Qwen36LoraTargetModule::SharedUpProj,
+                Qwen36LoraTargetModule::SharedDownProj,
+                Qwen36LoraTargetModule::ExpertsGateUpProj,
+                Qwen36LoraTargetModule::ExpertsDownProj,
+            ]
+        } else {
+            &[
+                Qwen36LoraTargetModule::GateProj,
+                Qwen36LoraTargetModule::UpProj,
+                Qwen36LoraTargetModule::DownProj,
+            ]
+        };
+        for &module in mlp_modules {
+            slots.push(Qwen36NativeLoraSlot {
+                index: slots.len(),
+                layer,
+                module,
+                active: (all_layers || target_layers.contains(&layer))
+                    && (all_modules || target_modules.contains(&module)),
+            });
+        }
+    }
+    slots
+}
+
+pub fn validate_lora_targets(
+    runtime_config: &Qwen36RuntimeConfig,
+    lora_config: &Qwen36LoraConfig,
+) -> Result<()> {
+    let slots = native_lora_slots(runtime_config, lora_config);
+    if !slots.iter().any(|slot| slot.active) {
+        bail!("LoRA targets do not resolve to any projection in this model");
+    }
+    if !runtime_config.is_moe
+        && lora_config.target_modules.iter().any(|module| {
+            matches!(
+                module,
+                Qwen36LoraTargetModule::SharedGateProj
+                    | Qwen36LoraTargetModule::SharedUpProj
+                    | Qwen36LoraTargetModule::SharedDownProj
+                    | Qwen36LoraTargetModule::ExpertsGateUpProj
+                    | Qwen36LoraTargetModule::ExpertsDownProj
+            )
+        })
+    {
+        bail!("shared expert LoRA targets require a MoE Qwen model");
+    }
+    if runtime_config.is_moe
+        && lora_config.target_modules.iter().any(|module| {
+            matches!(
+                module,
+                Qwen36LoraTargetModule::GateProj
+                    | Qwen36LoraTargetModule::UpProj
+                    | Qwen36LoraTargetModule::DownProj
+            )
+        })
+    {
+        bail!("dense MLP LoRA targets require a dense Qwen model");
+    }
+    Ok(())
+}
+
+fn adapter_tensor_prefix(
+    config: &Qwen36RuntimeConfig,
+    layer: usize,
+    module: Qwen36LoraTargetModule,
+) -> String {
+    format!(
+        "base_model.model.{}layers.{layer}.{}",
+        config.weight_prefix,
+        module.suffix()
+    )
+}
+
+/// Parse a canonical PEFT tensor name back to its layer, projection and side.
+pub fn parse_adapter_tensor_name(
+    config: &Qwen36RuntimeConfig,
+    name: &str,
+) -> Result<(usize, Qwen36LoraTargetModule, bool)> {
+    let (base, is_b) = if let Some(base) = name.strip_suffix(".lora_A.weight") {
+        (base, false)
+    } else if let Some(base) = name.strip_suffix(".lora_B.weight") {
+        (base, true)
+    } else {
+        bail!("invalid LoRA tensor name: {name}");
+    };
+    let prefix = format!("base_model.model.{}layers.", config.weight_prefix);
+    let rest = base
+        .strip_prefix(&prefix)
+        .with_context(|| format!("LoRA tensor has unexpected module prefix: {name}"))?;
+    let (layer, suffix) = rest
+        .split_once('.')
+        .with_context(|| format!("LoRA tensor is missing projection path: {name}"))?;
+    let layer = layer
+        .parse::<usize>()
+        .with_context(|| format!("invalid LoRA layer in tensor name: {name}"))?;
+    let module = [
+        Qwen36LoraTargetModule::QProj,
+        Qwen36LoraTargetModule::KProj,
+        Qwen36LoraTargetModule::VProj,
+        Qwen36LoraTargetModule::OProj,
+        Qwen36LoraTargetModule::InProjQkv,
+        Qwen36LoraTargetModule::InProjZ,
+        Qwen36LoraTargetModule::InProjA,
+        Qwen36LoraTargetModule::InProjB,
+        Qwen36LoraTargetModule::OutProj,
+        Qwen36LoraTargetModule::GateProj,
+        Qwen36LoraTargetModule::UpProj,
+        Qwen36LoraTargetModule::DownProj,
+        Qwen36LoraTargetModule::SharedGateProj,
+        Qwen36LoraTargetModule::SharedUpProj,
+        Qwen36LoraTargetModule::SharedDownProj,
+        Qwen36LoraTargetModule::ExpertsGateUpProj,
+        Qwen36LoraTargetModule::ExpertsDownProj,
+    ]
+    .into_iter()
+    .find(|module| module.suffix() == suffix)
+    .with_context(|| format!("unknown Qwen LoRA projection suffix: {suffix}"))?;
+    Ok((layer, module, is_b))
+}
+
+fn resolve_artifact_paths(path: &Path) -> (PathBuf, PathBuf) {
+    if path.extension().and_then(|extension| extension.to_str()) == Some("safetensors") {
+        let config_path = path
+            .parent()
+            .unwrap_or_else(|| Path::new("."))
+            .join("adapter_config.json");
+        (path.to_path_buf(), config_path)
+    } else {
+        (
+            path.join("adapter_model.safetensors"),
+            path.join("adapter_config.json"),
+        )
+    }
+}
+
+fn rustrain_config_path(config_path: &Path) -> PathBuf {
+    config_path
+        .parent()
+        .unwrap_or_else(|| Path::new("."))
+        .join("rustrain_adapter.json")
+}
+
+impl Qwen36AdapterArtifact {
+    pub fn from_native_exports(
+        model_name: &str,
+        architecture: &str,
+        base_model_path: Option<&Path>,
+        runtime_config: &Qwen36RuntimeConfig,
+        lora_config: &Qwen36LoraConfig,
+        exported: Vec<(Tensor, Tensor)>,
+    ) -> Result<Self> {
+        let slots = native_lora_slots(runtime_config, lora_config);
+        if exported.len() != slots.len() {
+            bail!(
+                "native LoRA export returned {} slots, expected {} for this model",
+                exported.len(),
+                slots.len()
+            );
+        }
+
+        let target_layers = slots
+            .iter()
+            .filter(|slot| slot.active)
+            .map(|slot| slot.layer)
+            .collect::<BTreeSet<_>>()
+            .into_iter()
+            .collect::<Vec<_>>();
+        let target_modules = slots
+            .iter()
+            .filter(|slot| slot.active)
+            .map(|slot| slot.module.cpp_name().to_string())
+            .collect::<BTreeSet<_>>()
+            .into_iter()
+            .collect::<Vec<_>>();
+
+        let mut tensors = BTreeMap::new();
+        for (slot, (a, b)) in slots.into_iter().zip(exported) {
+            if !slot.active {
+                continue;
+            }
+            let prefix = adapter_tensor_prefix(runtime_config, slot.layer, slot.module);
+            tensors.insert(
+                format!("{prefix}.lora_A.weight"),
+                a.to_device(tch::Device::Cpu).to_kind(Kind::Float),
+            );
+            tensors.insert(
+                format!("{prefix}.lora_B.weight"),
+                b.to_device(tch::Device::Cpu).to_kind(Kind::Float),
+            );
+        }
+        if tensors.is_empty() {
+            bail!("LoRA export contains no active target modules");
+        }
+
+        Ok(Self {
+            config: Qwen36AdapterConfig {
+                format_version: 1,
+                peft_type: "LORA".to_string(),
+                task_type: "CAUSAL_LM".to_string(),
+                base_model_name_or_path: base_model_path
+                    .map(|path| path.to_string_lossy().into_owned())
+                    .unwrap_or_else(|| model_name.to_string()),
+                rustrain_architecture: architecture.to_string(),
+                model_family: "qwen3_hybrid_text".to_string(),
+                r: lora_config.rank,
+                lora_alpha: lora_config.alpha,
+                target_layers,
+                target_modules,
+                adapter_dtype: "float32".to_string(),
+                bias: "none".to_string(),
+                inference_mode: true,
+            },
+            tensors,
+        })
+    }
+
+    pub fn save(&self, path: &Path) -> Result<PathBuf> {
+        let (tensor_path, config_path) = resolve_artifact_paths(path);
+        for (name, tensor) in &self.tensors {
+            if tensor.isfinite().all().int64_value(&[]) != 1 {
+                bail!("refusing to save non-finite LoRA tensor {name}");
+            }
+        }
+        if let Some(parent) = tensor_path.parent() {
+            fs::create_dir_all(parent)
+                .with_context(|| format!("failed to create {}", parent.display()))?;
+        }
+        let refs = self
+            .tensors
+            .iter()
+            .map(|(name, tensor)| (name.as_str(), tensor))
+            .collect::<Vec<_>>();
+        Tensor::write_safetensors(&refs, &tensor_path)
+            .with_context(|| format!("failed to write {}", tensor_path.display()))?;
+        let peft_config = PeftAdapterConfig {
+            peft_type: "LORA",
+            task_type: "CAUSAL_LM",
+            base_model_name_or_path: &self.config.base_model_name_or_path,
+            r: self.config.r,
+            lora_alpha: self.config.lora_alpha,
+            lora_dropout: 0.0,
+            target_modules: &self.config.target_modules,
+            layers_to_transform: &self.config.target_layers,
+            layers_pattern: "layers",
+            bias: "none",
+            fan_in_fan_out: false,
+            inference_mode: self.config.inference_mode,
+        };
+        fs::write(&config_path, serde_json::to_vec_pretty(&peft_config)?)
+            .with_context(|| format!("failed to write {}", config_path.display()))?;
+        let rustrain_path = rustrain_config_path(&config_path);
+        fs::write(&rustrain_path, serde_json::to_vec_pretty(&self.config)?)
+            .with_context(|| format!("failed to write {}", rustrain_path.display()))?;
+        Ok(tensor_path)
+    }
+
+    pub fn load(path: &Path) -> Result<Self> {
+        let (tensor_path, config_path) = resolve_artifact_paths(path);
+        let rustrain_path = rustrain_config_path(&config_path);
+        let metadata_path = if rustrain_path.exists() {
+            &rustrain_path
+        } else {
+            // Backward compatibility with the first named artifact draft,
+            // which stored rustrain metadata directly in adapter_config.json.
+            &config_path
+        };
+        let config: Qwen36AdapterConfig = if rustrain_path.exists() {
+            serde_json::from_slice(
+                &fs::read(metadata_path)
+                    .with_context(|| format!("failed to read {}", metadata_path.display()))?,
+            )
+            .with_context(|| format!("failed to parse {}", metadata_path.display()))?
+        } else {
+            // Standard PEFT exports do not carry rustrain-specific fields.
+            // Preserve their canonical metadata and fill only the native
+            // runtime fields that are not part of the PEFT schema.
+            let peft: PeftAdapterConfigOwned = serde_json::from_slice(
+                &fs::read(&config_path)
+                    .with_context(|| format!("failed to read {}", config_path.display()))?,
+            )
+            .with_context(|| format!("failed to parse {}", config_path.display()))?;
+            if peft.target_modules.is_empty() {
+                bail!("PEFT adapter_config.json has no target_modules");
+            }
+            Qwen36AdapterConfig {
+                format_version: 1,
+                peft_type: peft.peft_type,
+                task_type: peft.task_type,
+                base_model_name_or_path: peft.base_model_name_or_path,
+                rustrain_architecture: "qwen3_hybrid_lora_sft".to_string(),
+                model_family: "qwen3_hybrid_text".to_string(),
+                r: peft.r,
+                lora_alpha: peft.lora_alpha,
+                target_layers: peft.layers_to_transform,
+                target_modules: peft.target_modules,
+                adapter_dtype: "float32".to_string(),
+                bias: peft.bias,
+                inference_mode: peft.inference_mode,
+            }
+        };
+        if config.format_version != 1 || config.peft_type != "LORA" {
+            bail!(
+                "unsupported adapter format version/type: {}/{}",
+                config.format_version,
+                config.peft_type
+            );
+        }
+        if config.r <= 0 || config.lora_alpha <= 0.0 {
+            bail!("adapter rank and alpha must be positive");
+        }
+
+        let tensors = Tensor::read_safetensors(&tensor_path)
+            .with_context(|| format!("failed to read {}", tensor_path.display()))?
+            .into_iter()
+            .collect::<BTreeMap<_, _>>();
+        validate_adapter_tensors(&config, &tensors)?;
+        Ok(Self { config, tensors })
+    }
+
+    /// Convert the pre-registry positional export into the named v1 format.
+    /// No metadata is inferred: callers must provide the runtime target
+    /// contract that defined the positional slots.
+    pub fn load_legacy(
+        path: &Path,
+        model_name: &str,
+        architecture: &str,
+        runtime_config: &Qwen36RuntimeConfig,
+        lora_config: &Qwen36LoraConfig,
+    ) -> Result<Self> {
+        let positional = Tensor::read_safetensors(path)
+            .with_context(|| format!("failed to read legacy adapter {}", path.display()))?
+            .into_iter()
+            .collect::<BTreeMap<_, _>>();
+        let slots = native_lora_slots(runtime_config, lora_config);
+        let active_slots = slots.iter().filter(|slot| slot.active).collect::<Vec<_>>();
+        let a_count = positional
+            .keys()
+            .filter(|key| key.starts_with("lora_a_"))
+            .count();
+        let b_count = positional
+            .keys()
+            .filter(|key| key.starts_with("lora_b_"))
+            .count();
+        if a_count != b_count || (a_count != slots.len() && a_count != active_slots.len()) {
+            bail!(
+                "legacy adapter slot count {a_count}/{b_count} does not match fixed layout {} or compact target layout {}",
+                slots.len(),
+                active_slots.len()
+            );
+        }
+        let compact = a_count == active_slots.len() && a_count != slots.len();
+        let mut exported = Vec::with_capacity(slots.len());
+        let mut compact_index = 0usize;
+        for slot in &slots {
+            let source_index = if compact {
+                if !slot.active {
+                    let placeholder = Tensor::zeros([], (Kind::Float, tch::Device::Cpu));
+                    exported.push((placeholder.shallow_clone(), placeholder));
+                    continue;
+                }
+                let index = compact_index;
+                compact_index += 1;
+                index
+            } else {
+                slot.index
+            };
+            let a_name = format!("lora_a_{source_index}");
+            let b_name = format!("lora_b_{source_index}");
+            let a = positional
+                .get(&a_name)
+                .with_context(|| format!("legacy adapter missing {a_name}"))?
+                .shallow_clone();
+            let b = positional
+                .get(&b_name)
+                .with_context(|| format!("legacy adapter missing {b_name}"))?
+                .shallow_clone();
+            exported.push((a, b));
         }
+        Self::from_native_exports(
+            model_name,
+            architecture,
+            None,
+            runtime_config,
+            lora_config,
+            exported,
+        )
     }
 }
 
+fn validate_adapter_tensors(
+    config: &Qwen36AdapterConfig,
+    tensors: &BTreeMap<String, Tensor>,
+) -> Result<()> {
+    if tensors.is_empty() || tensors.len() % 2 != 0 {
+        bail!("adapter must contain paired LoRA A/B tensors");
+    }
+    let mut a_count = 0usize;
+    for (name, a) in tensors {
+        let Some(prefix) = name.strip_suffix(".lora_A.weight") else {
+            if name.ends_with(".lora_B.weight") {
+                let prefix = name.trim_end_matches(".lora_B.weight");
+                if !tensors.contains_key(&format!("{prefix}.lora_A.weight")) {
+                    bail!("missing paired tensor {prefix}.lora_A.weight");
+                }
+                continue;
+            }
+            bail!("unexpected adapter tensor name: {name}");
+        };
+        a_count += 1;
+        let b_name = format!("{prefix}.lora_B.weight");
+        let b = tensors
+            .get(&b_name)
+            .with_context(|| format!("missing paired tensor {b_name}"))?;
+        let grouped_expert = prefix.contains(".mlp.experts.");
+        if grouped_expert {
+            if a.dim() != 3 || b.dim() != 3 {
+                bail!("routed expert adapter tensors must be rank-3: {name}, {b_name}");
+            }
+            if a.size()[0] != b.size()[0] || a.size()[1] != config.r || b.size()[2] != config.r {
+                bail!(
+                    "routed expert adapter shape/rank mismatch for {prefix}: expected rank {}",
+                    config.r
+                );
+            }
+            if a.size()[2] <= 0 || b.size()[1] <= 0 {
+                bail!("routed expert adapter has an empty feature dimension: {prefix}");
+            }
+        } else {
+            if a.dim() != 2 || b.dim() != 2 {
+                bail!("adapter tensors must be matrices: {name}, {b_name}");
+            }
+            if a.size()[0] != config.r || b.size()[1] != config.r {
+                bail!("adapter rank mismatch for {prefix}: expected {}", config.r);
+            }
+            if a.size()[1] <= 0 || b.size()[0] <= 0 {
+                bail!("adapter tensor has an empty feature dimension: {prefix}");
+            }
+        }
+    }
+    if a_count == 0 {
+        bail!("adapter must contain at least one LoRA A tensor");
+    }
+    Ok(())
+}
+
 #[derive(Debug, Clone)]
 pub struct Qwen36LoraConfig {
     pub rank: i64,
@@ -84,16 +702,57 @@ impl Qwen36LoraRegistry {
             let layer_prefix = format!("{}layers.{}", config.weight_prefix, layer_idx);
             for &module in &lora_config.target_modules {
                 let weight_name = match module {
-                    Qwen36LoraTargetModule::QProj => format!("{layer_prefix}.self_attn.q_proj.weight"),
-                    Qwen36LoraTargetModule::KProj => format!("{layer_prefix}.self_attn.k_proj.weight"),
-                    Qwen36LoraTargetModule::VProj => format!("{layer_prefix}.self_attn.v_proj.weight"),
-                    Qwen36LoraTargetModule::OProj => format!("{layer_prefix}.self_attn.o_proj.weight"),
-                    Qwen36LoraTargetModule::InProjQkv => format!("{layer_prefix}.linear_attn.in_proj_qkv.weight"),
-                    Qwen36LoraTargetModule::InProjZ => format!("{layer_prefix}.linear_attn.in_proj_z.weight"),
-                    Qwen36LoraTargetModule::OutProj => format!("{layer_prefix}.linear_attn.out_proj.weight"),
-                    Qwen36LoraTargetModule::SharedGateProj => format!("{layer_prefix}.mlp.shared_expert.gate_proj.weight"),
-                    Qwen36LoraTargetModule::SharedUpProj => format!("{layer_prefix}.mlp.shared_expert.up_proj.weight"),
-                    Qwen36LoraTargetModule::SharedDownProj => format!("{layer_prefix}.mlp.shared_expert.down_proj.weight"),
+                    Qwen36LoraTargetModule::QProj => {
+                        format!("{layer_prefix}.self_attn.q_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::KProj => {
+                        format!("{layer_prefix}.self_attn.k_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::VProj => {
+                        format!("{layer_prefix}.self_attn.v_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::OProj => {
+                        format!("{layer_prefix}.self_attn.o_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::InProjQkv => {
+                        format!("{layer_prefix}.linear_attn.in_proj_qkv.weight")
+                    }
+                    Qwen36LoraTargetModule::InProjZ => {
+                        format!("{layer_prefix}.linear_attn.in_proj_z.weight")
+                    }
+                    Qwen36LoraTargetModule::InProjA => {
+                        format!("{layer_prefix}.linear_attn.in_proj_a.weight")
+                    }
+                    Qwen36LoraTargetModule::InProjB => {
+                        format!("{layer_prefix}.linear_attn.in_proj_b.weight")
+                    }
+                    Qwen36LoraTargetModule::OutProj => {
+                        format!("{layer_prefix}.linear_attn.out_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::GateProj => {
+                        format!("{layer_prefix}.mlp.gate_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::UpProj => {
+                        format!("{layer_prefix}.mlp.up_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::DownProj => {
+                        format!("{layer_prefix}.mlp.down_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::SharedGateProj => {
+                        format!("{layer_prefix}.mlp.shared_expert.gate_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::SharedUpProj => {
+                        format!("{layer_prefix}.mlp.shared_expert.up_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::SharedDownProj => {
+                        format!("{layer_prefix}.mlp.shared_expert.down_proj.weight")
+                    }
+                    Qwen36LoraTargetModule::ExpertsGateUpProj => {
+                        format!("{layer_prefix}.mlp.experts.gate_up_proj")
+                    }
+                    Qwen36LoraTargetModule::ExpertsDownProj => {
+                        format!("{layer_prefix}.mlp.experts.down_proj")
+                    }
                 };
 
                 let base_weight = match weights.get(&weight_name) {
@@ -101,34 +760,77 @@ impl Qwen36LoraRegistry {
                     None => {
                         // Skip modules that don't exist for this layer type
                         // (e.g., q_proj doesn't exist for linear attention layers)
-                        tracing::debug!("skipping LoRA target {weight_name} — not found (layer {layer_idx} may be different attention type)");
+                        tracing::debug!(
+                            "skipping LoRA target {weight_name} — not found (layer {layer_idx} may be different attention type)"
+                        );
                         continue;
                     }
                 };
 
-                let (out_features, in_features) = (base_weight.size()[0], base_weight.size()[1]);
                 let name = module.suffix().replace('.', "_");
                 let scale = 1.0 / (lora_config.rank as f64).sqrt();
-                let lora_a = p.randn(&format!("lora_a_{layer_idx}_{name}"), &[lora_config.rank, in_features], 0.0, scale);
-                let lora_b = p.zeros(&format!("lora_b_{layer_idx}_{name}"), &[out_features, lora_config.rank]);
+                let (lora_a, lora_b) = if base_weight.dim() == 3 {
+                    let experts = base_weight.size()[0];
+                    let out_features = base_weight.size()[1];
+                    let in_features = base_weight.size()[2];
+                    (
+                        p.randn(
+                            &format!("lora_a_{layer_idx}_{name}"),
+                            &[experts, lora_config.rank, in_features],
+                            0.0,
+                            scale,
+                        ),
+                        p.zeros(
+                            &format!("lora_b_{layer_idx}_{name}"),
+                            &[experts, out_features, lora_config.rank],
+                        ),
+                    )
+                } else {
+                    let out_features = base_weight.size()[0];
+                    let in_features = base_weight.size()[1];
+                    (
+                        p.randn(
+                            &format!("lora_a_{layer_idx}_{name}"),
+                            &[lora_config.rank, in_features],
+                            0.0,
+                            scale,
+                        ),
+                        p.zeros(
+                            &format!("lora_b_{layer_idx}_{name}"),
+                            &[out_features, lora_config.rank],
+                        ),
+                    )
+                };
                 adapters.insert((layer_idx, module), (lora_a, lora_b));
             }
         }
 
-        Ok(Self { config: lora_config, var_store, adapters })
+        Ok(Self {
+            config: lora_config,
+            var_store,
+            adapters,
+        })
     }
 
     pub fn trainable_variables(&self) -> Vec<Tensor> {
         self.var_store.trainable_variables()
     }
 
-    pub fn adapter_tensors(&self, layer: usize, module: Qwen36LoraTargetModule) -> Option<(Tensor, Tensor)> {
+    pub fn adapter_tensors(
+        &self,
+        layer: usize,
+        module: Qwen36LoraTargetModule,
+    ) -> Option<(Tensor, Tensor)> {
         let (a, b) = self.adapters.get(&(layer, module))?;
         Some((a.shallow_clone(), b.shallow_clone()))
     }
 
     /// Get references to the actual VarStore adapter tensors (for training backward).
-    pub fn adapter_ref(&self, layer: usize, module: Qwen36LoraTargetModule) -> Option<(&Tensor, &Tensor)> {
+    pub fn adapter_ref(
+        &self,
+        layer: usize,
+        module: Qwen36LoraTargetModule,
+    ) -> Option<(&Tensor, &Tensor)> {
         let (a, b) = self.adapters.get(&(layer, module))?;
         Some((a, b))
     }
@@ -138,7 +840,10 @@ impl Qwen36LoraRegistry {
     }
 
     pub fn trainable_param_count(&self) -> usize {
-        self.trainable_variables().iter().map(|v| v.numel() as usize).sum()
+        self.trainable_variables()
+            .iter()
+            .map(|v| v.numel() as usize)
+            .sum()
     }
 
     pub fn save(&self, path: &std::path::Path) -> Result<()> {
@@ -156,16 +861,22 @@ impl Qwen36LoraRegistry {
         let mut header = serde_json::Map::new();
         let mut offset = 0u64;
         for (name, tensor) in &named {
-            let t = tensor.to_device(tch::Device::Cpu).contiguous().to_kind(Kind::Float);
+            let t = tensor
+                .to_device(tch::Device::Cpu)
+                .contiguous()
+                .to_kind(Kind::Float);
             let shape: Vec<i64> = t.size().iter().copied().map(|d| d).collect();
             let t_flat = t.reshape([-1]);
             let data: Vec<f32> = Vec::<f32>::try_from(&t_flat)?;
             let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
-            header.insert(name.clone(), serde_json::json!({
-                "dtype": "F32",
-                "shape": shape,
-                "data_offsets": [offset, offset + bytes.len() as u64],
-            }));
+            header.insert(
+                name.clone(),
+                serde_json::json!({
+                    "dtype": "F32",
+                    "shape": shape,
+                    "data_offsets": [offset, offset + bytes.len() as u64],
+                }),
+            );
             offset += bytes.len() as u64;
             tensors_data.push(bytes);
         }
@@ -185,3 +896,141 @@ impl Qwen36LoraRegistry {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn runtime(prefix: &str, layer_types: Vec<LayerType>) -> Qwen36RuntimeConfig {
+        Qwen36RuntimeConfig {
+            num_hidden_layers: layer_types.len(),
+            hidden_size: 16,
+            vocab_size: 32,
+            rms_norm_eps: 1e-6,
+            tie_word_embeddings: true,
+            hidden_act: "silu".into(),
+            layer_types,
+            full_attention_interval: 4,
+            num_attention_heads: 2,
+            num_key_value_heads: 2,
+            head_dim: 8,
+            attention_bias: false,
+            attn_output_gate: false,
+            rope_theta: 1e6,
+            partial_rotary_factor: 1.0,
+            mrope_interleaved: false,
+            mrope_section: vec![],
+            linear_num_key_heads: 2,
+            linear_key_head_dim: 8,
+            linear_num_value_heads: 2,
+            linear_value_head_dim: 8,
+            linear_conv_kernel_dim: 4,
+            mamba_ssm_dtype: "float32".into(),
+            is_moe: false,
+            num_experts: 0,
+            num_experts_per_tok: 0,
+            moe_intermediate_size: 0,
+            shared_expert_intermediate_size: 0,
+            norm_topk_prob: true,
+            router_aux_loss_coef: 0.0,
+            intermediate_size: 32,
+            mtp_num_hidden_layers: 0,
+            mtp_use_dedicated_embeddings: false,
+            has_vision: prefix.contains("language_model"),
+            vision_depth: 0,
+            vision_hidden_size: 0,
+            vision_num_heads: 0,
+            vision_patch_size: 0,
+            vision_spatial_merge_size: 0,
+            vision_temporal_patch_size: 0,
+            vision_out_hidden_size: 0,
+            weight_prefix: prefix.into(),
+        }
+    }
+
+    #[test]
+    fn native_slots_match_cpp_projection_order() {
+        let config = runtime(
+            "model.",
+            vec![LayerType::FullAttention, LayerType::LinearAttention],
+        );
+        let lora = Qwen36LoraConfig {
+            rank: 2,
+            alpha: 4.0,
+            target_layers: vec![],
+            target_modules: vec![],
+        };
+        let slots = native_lora_slots(&config, &lora);
+        let modules = slots
+            .iter()
+            .map(|slot| slot.module.cpp_name())
+            .collect::<Vec<_>>();
+        assert_eq!(
+            modules,
+            [
+                "q_proj",
+                "k_proj",
+                "v_proj",
+                "o_proj",
+                "gate_proj",
+                "up_proj",
+                "down_proj",
+                "in_proj_qkv",
+                "in_proj_z",
+                "in_proj_a",
+                "in_proj_b",
+                "out_proj",
+                "gate_proj",
+                "up_proj",
+                "down_proj"
+            ]
+        );
+    }
+
+    #[test]
+    fn parses_text_and_multimodal_peft_names() {
+        let text = runtime("model.", vec![LayerType::FullAttention]);
+        assert_eq!(
+            parse_adapter_tensor_name(
+                &text,
+                "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight"
+            )
+            .unwrap(),
+            (0, Qwen36LoraTargetModule::QProj, false)
+        );
+        let multimodal = runtime("model.language_model.", vec![LayerType::LinearAttention]);
+        assert_eq!(
+            parse_adapter_tensor_name(
+                &multimodal,
+                "base_model.model.model.language_model.layers.0.linear_attn.in_proj_qkv.lora_B.weight"
+            )
+            .unwrap(),
+            (0, Qwen36LoraTargetModule::InProjQkv, true)
+        );
+    }
+
+    #[test]
+    fn rejects_unpaired_adapter_tensor() {
+        let config = Qwen36AdapterConfig {
+            format_version: 1,
+            peft_type: "LORA".into(),
+            task_type: "CAUSAL_LM".into(),
+            base_model_name_or_path: "Qwen/test".into(),
+            rustrain_architecture: "qwen3.6".into(),
+            model_family: "qwen3_hybrid_text".into(),
+            r: 2,
+            lora_alpha: 4.0,
+            target_layers: vec![0],
+            target_modules: vec!["q_proj".into()],
+            adapter_dtype: "float32".into(),
+            bias: "none".into(),
+            inference_mode: true,
+        };
+        let mut tensors = BTreeMap::new();
+        tensors.insert(
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight".into(),
+            Tensor::zeros([4, 2], (Kind::Float, tch::Device::Cpu)),
+        );
+        assert!(validate_adapter_tensors(&config, &tensors).is_err());
+    }
+}
diff --git a/crates/rustrain-qwen3-6/src/model.rs b/crates/rustrain-qwen3-6/src/model.rs
index c3f01f0c..1fe436cf 100644
--- a/crates/rustrain-qwen3-6/src/model.rs
+++ b/crates/rustrain-qwen3-6/src/model.rs
@@ -316,10 +316,12 @@ pub fn linear_attention(
     let k = k.repeat_interleave_self_int(n_rep, 2, None);
 
     // L2 normalize Q, K (HF uses use_qk_l2norm_in_kernel=True)
-    let q_norm = q.pow_tensor_scalar(2.0).sum_dim_intlist([-1].as_slice(), true, compute_kind).sqrt().clamp_min(1e-6);
-    let k_norm = k.pow_tensor_scalar(2.0).sum_dim_intlist([-1].as_slice(), true, compute_kind).sqrt().clamp_min(1e-6);
-    let q = (&q / &q_norm).to_kind(Kind::Float);
-    let k = (&k / &k_norm).to_kind(Kind::Float);
+    let q = q.to_kind(Kind::Float);
+    let k = k.to_kind(Kind::Float);
+    let q_norm = (q.pow_tensor_scalar(2.0).sum_dim_intlist([-1].as_slice(), true, Kind::Float) + 1e-6).rsqrt();
+    let k_norm = (k.pow_tensor_scalar(2.0).sum_dim_intlist([-1].as_slice(), true, Kind::Float) + 1e-6).rsqrt();
+    let q = &q * &q_norm;
+    let k = &k * &k_norm;
 
     // Scale Q by 1/sqrt(key_dim) — matching HF: scale = 1 / (query.shape[-1] ** 0.5)
     let scale = 1.0 / (key_dim as f64).sqrt();
diff --git a/crates/rustrain-qwen3-6/src/pipeline.rs b/crates/rustrain-qwen3-6/src/pipeline.rs
new file mode 100644
index 00000000..cea900bc
--- /dev/null
+++ b/crates/rustrain-qwen3-6/src/pipeline.rs
@@ -0,0 +1,321 @@
+use std::collections::HashSet;
+use std::ops::Range;
+
+use anyhow::{Result, bail};
+
+use crate::config::{LayerType, Qwen36RuntimeConfig};
+use crate::lora::{Qwen36LoraTargetModule, Qwen36NativeLoraSlot};
+
+#[derive(Debug, Clone)]
+pub struct PipelineLoraSlot {
+    pub global_index: usize,
+    pub local_index: usize,
+    pub layer: usize,
+    pub module: Qwen36LoraTargetModule,
+    pub active: bool,
+}
+
+/// Contiguous pipeline ownership for one physical PP rank.
+///
+/// Layer IDs remain global at every external boundary. `layer_range` is the
+/// only place where they become a stage-local contiguous slice.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct PipelineStageLayout {
+    pub pipeline_rank: usize,
+    pub pipeline_size: usize,
+    pub global_num_layers: usize,
+    pub layer_range: Range<usize>,
+}
+
+impl PipelineStageLayout {
+    pub fn new(
+        global_num_layers: usize,
+        pipeline_rank: usize,
+        pipeline_size: usize,
+    ) -> Result<Self> {
+        if pipeline_size == 0 {
+            bail!("pipeline_size must be positive");
+        }
+        if pipeline_rank >= pipeline_size {
+            bail!("pipeline rank {pipeline_rank} is outside pipeline_size={pipeline_size}");
+        }
+        if global_num_layers < pipeline_size {
+            bail!(
+                "global_num_layers={global_num_layers} must be at least pipeline_size={pipeline_size}"
+            );
+        }
+        let start = global_num_layers * pipeline_rank / pipeline_size;
+        let end = global_num_layers * (pipeline_rank + 1) / pipeline_size;
+        Ok(Self {
+            pipeline_rank,
+            pipeline_size,
+            global_num_layers,
+            layer_range: start..end,
+        })
+    }
+
+    pub fn full(global_num_layers: usize) -> Result<Self> {
+        Self::new(global_num_layers, 0, 1)
+    }
+
+    pub fn is_first(&self) -> bool {
+        self.pipeline_rank == 0
+    }
+
+    pub fn is_last(&self) -> bool {
+        self.pipeline_rank + 1 == self.pipeline_size
+    }
+
+    pub fn local_num_layers(&self) -> usize {
+        self.layer_range.len()
+    }
+
+    pub fn owns_layer(&self, global_layer: usize) -> bool {
+        self.layer_range.contains(&global_layer)
+    }
+
+    pub fn local_target_layers(&self, global_targets: &[usize]) -> Vec<usize> {
+        global_targets
+            .iter()
+            .copied()
+            .filter(|layer| self.owns_layer(*layer))
+            .collect()
+    }
+
+    pub(crate) fn native_flags(&self) -> i32 {
+        i32::from(self.is_first()) | (i32::from(self.is_last()) << 1)
+    }
+}
+
+pub fn stage_lora_slots(
+    global_slots: &[Qwen36NativeLoraSlot],
+    stage: &PipelineStageLayout,
+) -> Vec<PipelineLoraSlot> {
+    global_slots
+        .iter()
+        .filter(|slot| stage.owns_layer(slot.layer))
+        .enumerate()
+        .map(|(local_index, slot)| PipelineLoraSlot {
+            global_index: slot.index,
+            local_index,
+            layer: slot.layer,
+            module: slot.module,
+            active: slot.active,
+        })
+        .collect()
+}
+
+/// Frozen base weights owned by one pipeline stage.
+///
+/// Tied vocabulary weights are intentionally present on both boundary stages:
+/// the first consumes them as embeddings and the last consumes them as the LM
+/// head. All other layer weights have exactly one owner.
+pub fn stage_text_needed_weights(
+    config: &Qwen36RuntimeConfig,
+    stage: &PipelineStageLayout,
+) -> HashSet<String> {
+    let prefix = &config.weight_prefix;
+    let mut needed = HashSet::new();
+
+    if stage.is_first() || (stage.is_last() && config.tie_word_embeddings) {
+        needed.insert(format!("{prefix}embed_tokens.weight"));
+    }
+    if stage.is_last() {
+        needed.insert(format!("{prefix}norm.weight"));
+        if !config.tie_word_embeddings {
+            needed.insert("lm_head.weight".to_string());
+        }
+    }
+
+    for layer in stage.layer_range.clone() {
+        let layer_prefix = format!("{prefix}layers.{layer}");
+        needed.insert(format!("{layer_prefix}.input_layernorm.weight"));
+        needed.insert(format!("{layer_prefix}.post_attention_layernorm.weight"));
+
+        match config.layer_types[layer] {
+            LayerType::FullAttention => {
+                for weight in ["q_proj", "q_norm", "k_proj", "k_norm", "v_proj", "o_proj"] {
+                    needed.insert(format!("{layer_prefix}.self_attn.{weight}.weight"));
+                }
+            }
+            LayerType::LinearAttention => {
+                needed.insert(format!("{layer_prefix}.linear_attn.A_log"));
+                needed.insert(format!("{layer_prefix}.linear_attn.conv1d.weight"));
+                needed.insert(format!("{layer_prefix}.linear_attn.dt_bias"));
+                needed.insert(format!("{layer_prefix}.linear_attn.norm.weight"));
+                for weight in [
+                    "in_proj_qkv",
+                    "in_proj_z",
+                    "in_proj_a",
+                    "in_proj_b",
+                    "out_proj",
+                ] {
+                    needed.insert(format!("{layer_prefix}.linear_attn.{weight}.weight"));
+                }
+            }
+        }
+
+        if config.is_moe {
+            needed.insert(format!("{layer_prefix}.mlp.gate.weight"));
+            needed.insert(format!("{layer_prefix}.mlp.shared_expert_gate.weight"));
+            needed.insert(format!("{layer_prefix}.mlp.shared_expert.gate_proj.weight"));
+            needed.insert(format!("{layer_prefix}.mlp.shared_expert.up_proj.weight"));
+            needed.insert(format!("{layer_prefix}.mlp.shared_expert.down_proj.weight"));
+            needed.insert(format!("{layer_prefix}.mlp.experts.gate_up_proj"));
+            needed.insert(format!("{layer_prefix}.mlp.experts.down_proj"));
+        } else {
+            needed.insert(format!("{layer_prefix}.mlp.gate_proj.weight"));
+            needed.insert(format!("{layer_prefix}.mlp.up_proj.weight"));
+            needed.insert(format!("{layer_prefix}.mlp.down_proj.weight"));
+        }
+    }
+
+    needed
+}
+
+pub fn stage_needed_weights(
+    config: &Qwen36RuntimeConfig,
+    stage: &PipelineStageLayout,
+) -> HashSet<String> {
+    let mut needed = stage_text_needed_weights(config, stage);
+    if config.has_vision && stage.is_first() {
+        needed.extend(crate::vision::VisionWeights::weight_names(config));
+    }
+
+    needed
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{PipelineStageLayout, stage_lora_slots, stage_needed_weights};
+    use crate::config::{LayerType, Qwen36RuntimeConfig};
+    use crate::lora::{Qwen36LoraTargetModule, Qwen36NativeLoraSlot};
+
+    fn runtime() -> Qwen36RuntimeConfig {
+        Qwen36RuntimeConfig {
+            num_hidden_layers: 2,
+            hidden_size: 16,
+            vocab_size: 32,
+            rms_norm_eps: 1e-6,
+            tie_word_embeddings: true,
+            hidden_act: "silu".into(),
+            layer_types: vec![LayerType::FullAttention, LayerType::LinearAttention],
+            full_attention_interval: 2,
+            num_attention_heads: 2,
+            num_key_value_heads: 2,
+            head_dim: 8,
+            attention_bias: false,
+            attn_output_gate: false,
+            rope_theta: 1e6,
+            partial_rotary_factor: 1.0,
+            mrope_interleaved: false,
+            mrope_section: vec![],
+            linear_num_key_heads: 2,
+            linear_key_head_dim: 8,
+            linear_num_value_heads: 2,
+            linear_value_head_dim: 8,
+            linear_conv_kernel_dim: 4,
+            mamba_ssm_dtype: "float32".into(),
+            is_moe: false,
+            num_experts: 0,
+            num_experts_per_tok: 0,
+            moe_intermediate_size: 0,
+            shared_expert_intermediate_size: 0,
+            norm_topk_prob: true,
+            router_aux_loss_coef: 0.0,
+            intermediate_size: 32,
+            mtp_num_hidden_layers: 0,
+            mtp_use_dedicated_embeddings: false,
+            has_vision: false,
+            vision_depth: 0,
+            vision_hidden_size: 0,
+            vision_num_heads: 0,
+            vision_patch_size: 0,
+            vision_spatial_merge_size: 0,
+            vision_temporal_patch_size: 0,
+            vision_out_hidden_size: 0,
+            weight_prefix: "model.".into(),
+        }
+    }
+
+    #[test]
+    fn uneven_layers_are_contiguous_and_cover_the_model() {
+        let stages = (0..2)
+            .map(|rank| PipelineStageLayout::new(5, rank, 2).unwrap())
+            .collect::<Vec<_>>();
+        assert_eq!(stages[0].layer_range, 0..2);
+        assert_eq!(stages[1].layer_range, 2..5);
+        assert!(stages[0].is_first());
+        assert!(stages[1].is_last());
+        assert_eq!(stages[0].layer_range.end, stages[1].layer_range.start);
+        assert_eq!(stages[1].layer_range.end, 5);
+    }
+
+    #[test]
+    fn target_layers_keep_global_identity() {
+        let first = PipelineStageLayout::new(4, 0, 2).unwrap();
+        let last = PipelineStageLayout::new(4, 1, 2).unwrap();
+        assert_eq!(first.local_target_layers(&[0, 3]), vec![0]);
+        assert_eq!(last.local_target_layers(&[0, 3]), vec![3]);
+    }
+
+    #[test]
+    fn rejects_empty_stages_and_invalid_rank() {
+        assert!(PipelineStageLayout::new(1, 0, 2).is_err());
+        assert!(PipelineStageLayout::new(4, 2, 2).is_err());
+        assert!(PipelineStageLayout::new(4, 0, 0).is_err());
+    }
+
+    #[test]
+    fn stage_slots_preserve_global_identity_and_compact_native_indices() {
+        let global_slots = vec![
+            Qwen36NativeLoraSlot {
+                index: 0,
+                layer: 0,
+                module: Qwen36LoraTargetModule::QProj,
+                active: true,
+            },
+            Qwen36NativeLoraSlot {
+                index: 1,
+                layer: 0,
+                module: Qwen36LoraTargetModule::OProj,
+                active: false,
+            },
+            Qwen36NativeLoraSlot {
+                index: 2,
+                layer: 1,
+                module: Qwen36LoraTargetModule::InProjQkv,
+                active: true,
+            },
+        ];
+        let stage = PipelineStageLayout::new(2, 1, 2).unwrap();
+        let local = stage_lora_slots(&global_slots, &stage);
+        assert_eq!(local.len(), 1);
+        assert_eq!(local[0].global_index, 2);
+        assert_eq!(local[0].local_index, 0);
+        assert_eq!(local[0].layer, 1);
+    }
+
+    #[test]
+    fn boundary_weights_are_owned_without_loading_remote_layers() {
+        let config = runtime();
+        let first = stage_needed_weights(
+            &config,
+            &PipelineStageLayout::new(config.num_hidden_layers, 0, 2).unwrap(),
+        );
+        let last = stage_needed_weights(
+            &config,
+            &PipelineStageLayout::new(config.num_hidden_layers, 1, 2).unwrap(),
+        );
+
+        assert!(first.contains("model.embed_tokens.weight"));
+        assert!(!first.contains("model.norm.weight"));
+        assert!(first.iter().any(|name| name.starts_with("model.layers.0.")));
+        assert!(!first.iter().any(|name| name.starts_with("model.layers.1.")));
+
+        assert!(last.contains("model.embed_tokens.weight"));
+        assert!(last.contains("model.norm.weight"));
+        assert!(last.iter().any(|name| name.starts_with("model.layers.1.")));
+        assert!(!last.iter().any(|name| name.starts_with("model.layers.0.")));
+    }
+}
diff --git a/crates/rustrain-qwen3-6/src/session.rs b/crates/rustrain-qwen3-6/src/session.rs
index ee2a90ab..be0a337e 100644
--- a/crates/rustrain-qwen3-6/src/session.rs
+++ b/crates/rustrain-qwen3-6/src/session.rs
@@ -7,11 +7,19 @@ use anyhow::{Context, Result, anyhow, bail};
 use tch::{Kind, Tensor};
 use tracing::info;
 
-use crate::config::{read_qwen36_runtime_config, resolve_qwen36_model_path, Qwen36RuntimeConfig, LayerType};
-use crate::lora::{Qwen36LoraConfig, Qwen36LoraTargetModule};
+use crate::checkpoint;
+use crate::config::{
+    LayerType, Qwen36RuntimeConfig, read_qwen36_runtime_config, resolve_qwen36_model_path,
+};
+use crate::lora::{
+    Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule, native_lora_slots,
+    validate_lora_targets,
+};
+use crate::pipeline::{PipelineStageLayout, stage_lora_slots, stage_needed_weights};
 use crate::sft::SftDataset;
+use rustrain_checkpoint::safetensors::read_safetensors_dir_filtered;
 use rustrain_core::runtime::{Config, RunPaths};
-use rustrain_checkpoint::safetensors::{read_safetensors_dir_filtered};
+use rustrain_parallel::topology::{DEFAULT_RANK_ORDER, ParallelTopology};
 
 // ──────────────────────────────────────────────────────────────────────
 // EP Shard
@@ -27,7 +35,10 @@ pub struct EpShard {
 
 impl EpShard {
     pub fn new(rank: usize, world_size: usize, num_experts: usize) -> Self {
-        assert!(num_experts % world_size == 0, "num_experts {num_experts} not divisible by world_size {world_size}");
+        assert!(
+            num_experts % world_size == 0,
+            "num_experts {num_experts} not divisible by world_size {world_size}"
+        );
         let epr = num_experts / world_size;
         let start = rank * epr;
         Self {
@@ -62,11 +73,33 @@ pub struct Qwen36LoraSftSummary {
 fn parse_env_usize(key: &str) -> Result<usize> {
     env::var(key)
         .with_context(|| format!("{key} not set"))
-        .and_then(|v| v.parse::<usize>().with_context(|| format!("invalid {key}: {v}")))
+        .and_then(|v| {
+            v.parse::<usize>()
+                .with_context(|| format!("invalid {key}: {v}"))
+        })
+}
+
+fn validate_lora_rank_for_tp(
+    lora_rank: i64,
+    tp_size: usize,
+    layouts: &[checkpoint::LoraTpShardLayout],
+) -> Result<()> {
+    if tp_size > 1
+        && layouts
+            .iter()
+            .any(|layout| *layout == checkpoint::LoraTpShardLayout::LatentRank)
+        && lora_rank % tp_size as i64 != 0
+    {
+        bail!("latent-rank LoRA rank {lora_rank} must be divisible by TP_SIZE={tp_size}");
+    }
+    Ok(())
 }
 
 fn lora_config_from_config(config: &Config) -> Result<Qwen36LoraConfig> {
-    let lora = config.lora.as_ref().ok_or_else(|| anyhow!("[lora] section required"))?;
+    let lora = config
+        .lora
+        .as_ref()
+        .ok_or_else(|| anyhow!("[lora] section required"))?;
     let target_layers: Vec<usize> = lora.target_layers.clone();
     let target_modules: Vec<Qwen36LoraTargetModule> = lora
         .target_modules
@@ -81,73 +114,65 @@ fn lora_config_from_config(config: &Config) -> Result<Qwen36LoraConfig> {
     })
 }
 
-/// Build the set of weight names needed for training.
-/// For EP mode, only load local expert slice.
-fn build_needed_weights(
-    config: &Qwen36RuntimeConfig,
-    lora_config: &Qwen36LoraConfig,
-    ep_shard: Option<&EpShard>,
-) -> HashSet<String> {
-    let p = &config.weight_prefix;
-    let mut needed = HashSet::new();
-
-    // Embed, norm, lm_head
-    needed.insert(format!("{p}embed_tokens.weight"));
-    needed.insert(format!("{p}norm.weight"));
-    if !config.tie_word_embeddings {
-        needed.insert("lm_head.weight".to_string());
-    }
-
-    // Per-layer weights for ALL layers (not just LoRA targets)
-    for layer in 0..config.num_hidden_layers {
-        let lp = format!("{p}layers.{layer}");
-        needed.insert(format!("{lp}.input_layernorm.weight"));
-        needed.insert(format!("{lp}.post_attention_layernorm.weight"));
-
-        // Attention weights (full or linear)
-        match config.layer_types[layer] {
-            LayerType::FullAttention => {
-                for w in &["q_proj", "q_norm", "k_proj", "k_norm", "v_proj", "o_proj"] {
-                    needed.insert(format!("{lp}.self_attn.{w}.weight"));
-                }
-            }
-            LayerType::LinearAttention => {
-                needed.insert(format!("{lp}.linear_attn.A_log"));
-                needed.insert(format!("{lp}.linear_attn.conv1d.weight"));
-                needed.insert(format!("{lp}.linear_attn.dt_bias"));
-                needed.insert(format!("{lp}.linear_attn.norm.weight"));
-                for w in &["in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b", "out_proj"] {
-                    needed.insert(format!("{lp}.linear_attn.{w}.weight"));
-                }
-            }
-        }
+fn training_source_coordinate(
+    dp_rank: usize,
+    dp_size: usize,
+    ep_rank: usize,
+    ep_size: usize,
+    ep_source_sharded: bool,
+) -> (usize, usize) {
+    if ep_source_sharded {
+        (dp_rank * ep_size + ep_rank, dp_size * ep_size)
+    } else if dp_size > 1 {
+        (dp_rank, dp_size)
+    } else {
+        (0, 1)
+    }
+}
 
-        // MLP weights — dense vs MoE
-        if config.is_moe {
-            needed.insert(format!("{lp}.mlp.gate.weight"));
-            needed.insert(format!("{lp}.mlp.shared_expert_gate.weight"));
-            needed.insert(format!("{lp}.mlp.shared_expert.gate_proj.weight"));
-            needed.insert(format!("{lp}.mlp.shared_expert.up_proj.weight"));
-            needed.insert(format!("{lp}.mlp.shared_expert.down_proj.weight"));
-            // Fused expert tensors are 3D [num_experts, ...], loaded as a whole
-            needed.insert(format!("{lp}.mlp.experts.gate_up_proj"));
-            needed.insert(format!("{lp}.mlp.experts.down_proj"));
-        } else {
-            // Dense MLP: standard SwiGLU (gate_proj, up_proj, down_proj)
-            needed.insert(format!("{lp}.mlp.gate_proj.weight"));
-            needed.insert(format!("{lp}.mlp.up_proj.weight"));
-            needed.insert(format!("{lp}.mlp.down_proj.weight"));
-        }
+fn pipeline_1f1b_schedule(
+    pp_rank: usize,
+    pp_size: usize,
+    num_microbatches: usize,
+) -> Result<Vec<(Option<i64>, Option<i64>)>> {
+    if pp_size < 2 {
+        bail!("pipeline schedule requires PP_SIZE >= 2");
+    }
+    if pp_rank >= pp_size {
+        bail!("pipeline rank {pp_rank} is outside PP_SIZE={pp_size}");
+    }
+    if num_microbatches == 0 {
+        bail!("pipeline schedule requires at least one microbatch");
+    }
+    let warmup = (pp_size - pp_rank - 1).min(num_microbatches);
+    let mut schedule = Vec::with_capacity(num_microbatches + warmup);
+    for microbatch in 0..warmup {
+        schedule.push((
+            Some(i64::try_from(microbatch).context("microbatch id exceeds i64")?),
+            None,
+        ));
+    }
+    for microbatch in warmup..num_microbatches {
+        let forward = i64::try_from(microbatch).context("microbatch id exceeds i64")?;
+        let backward = i64::try_from(microbatch - warmup).context("microbatch id exceeds i64")?;
+        schedule.push((Some(forward), Some(backward)));
     }
+    for microbatch in num_microbatches - warmup..num_microbatches {
+        schedule.push((
+            None,
+            Some(i64::try_from(microbatch).context("microbatch id exceeds i64")?),
+        ));
+    }
+    Ok(schedule)
+}
 
-    // Vision encoder (for multimodal)
-    if config.has_vision {
-        for name in crate::vision::VisionWeights::weight_names(config) {
-            needed.insert(name);
+fn validate_native_optimizer_options(max_grad_norm: Option<f32>) -> Result<()> {
+    if let Some(max_grad_norm) = max_grad_norm {
+        if !max_grad_norm.is_finite() || max_grad_norm <= 0.0 {
+            bail!("train.max_grad_norm must be finite and greater than zero");
         }
     }
-
-    needed
+    Ok(())
 }
 
 // ──────────────────────────────────────────────────────────────────────
@@ -171,19 +196,154 @@ pub fn train_qwen3_6_lora_sft_ep(
 ) -> Result<Qwen36LoraSftSummary> {
     let rank = parse_env_usize("RANK")?;
     let world_size = parse_env_usize("WORLD_SIZE")?;
-    // For MoE models, shard experts. For dense models, EP is a no-op (no experts to shard).
-    let model_path = config.model.model_path.as_ref()
+    let model_path = config
+        .model
+        .model_path
+        .as_ref()
         .ok_or_else(|| anyhow!("model.model_path required"))?;
-    let model_path = resolve_qwen36_model_path(model_path)?;
+    let model_path = std::fs::canonicalize(resolve_qwen36_model_path(model_path)?)
+        .with_context(|| format!("canonicalize Qwen model path {}", model_path.display()))?;
     let runtime_config = read_qwen36_runtime_config(&model_path)?;
     let ep_shard = if runtime_config.is_moe {
-        Some(EpShard::new(rank, world_size, runtime_config.num_experts))
+        if config.parallel.context_parallel_size != 1 {
+            bail!(
+                "native MoE Qwen LoRA does not yet support context parallelism: TP={} PP={} DP={} EP={} CP={}",
+                config.parallel.tensor_model_parallel_size,
+                config.parallel.pipeline_model_parallel_size,
+                config.parallel.data_parallel_size,
+                config.parallel.expert_model_parallel_size,
+                config.parallel.context_parallel_size,
+            );
+        }
+        let rank_order = std::env::var("RUSTRAIN_PARALLEL_ORDER")
+            .or_else(|_| std::env::var("PARALLEL_ORDER"))
+            .unwrap_or_else(|_| DEFAULT_RANK_ORDER.to_string());
+        let topology = ParallelTopology::with_order(
+            config.parallel.tensor_model_parallel_size,
+            config.parallel.pipeline_model_parallel_size,
+            config.parallel.data_parallel_size,
+            config.parallel.expert_model_parallel_size,
+            1,
+            &rank_order,
+        )?;
+        topology.validate_world_size(world_size)?;
+        Some(EpShard::new(
+            topology.expert_rank(rank)?,
+            topology.expert_model_parallel_size(),
+            runtime_config.num_experts,
+        ))
     } else {
         None
     };
     train_impl(config, run_paths, ep_shard)
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn tp_peers_share_expert_and_source_shards() {
+        let topology = ParallelTopology::new(2, 1, 1, 2, 1).unwrap();
+        let expected_ep_ranks = [0, 0, 1, 1];
+        for (global_rank, expected_ep_rank) in expected_ep_ranks.into_iter().enumerate() {
+            let ep_rank = topology.expert_rank(global_rank).unwrap();
+            assert_eq!(ep_rank, expected_ep_rank);
+            let shard = EpShard::new(ep_rank, 2, 8);
+            assert_eq!(shard.expert_start, expected_ep_rank * 4);
+            let (source_rank, source_count) = training_source_coordinate(0, 1, ep_rank, 2, true);
+            let data_start = 3 * 2 * source_count + source_rank * 2;
+            assert_eq!(data_start, 12 + expected_ep_rank * 2);
+        }
+    }
+
+    #[test]
+    fn tp_ep_dp_uses_every_expert_and_data_source_coordinate() {
+        let expected = [(0, 4), (1, 4), (2, 4), (3, 4)];
+        let observed = [(0, 0), (0, 1), (1, 0), (1, 1)]
+            .map(|(dp_rank, ep_rank)| training_source_coordinate(dp_rank, 2, ep_rank, 2, true));
+        assert_eq!(observed, expected);
+        assert_eq!(training_source_coordinate(1, 2, 1, 2, false), (1, 2));
+    }
+
+    #[test]
+    fn pipeline_schedule_matches_non_interleaved_1f1b() {
+        assert_eq!(
+            pipeline_1f1b_schedule(0, 2, 4).unwrap(),
+            vec![
+                (Some(0), None),
+                (Some(1), Some(0)),
+                (Some(2), Some(1)),
+                (Some(3), Some(2)),
+                (None, Some(3)),
+            ]
+        );
+        assert_eq!(
+            pipeline_1f1b_schedule(1, 2, 4).unwrap(),
+            vec![
+                (Some(0), Some(0)),
+                (Some(1), Some(1)),
+                (Some(2), Some(2)),
+                (Some(3), Some(3)),
+            ]
+        );
+        assert_eq!(
+            pipeline_1f1b_schedule(1, 3, 4).unwrap(),
+            vec![
+                (Some(0), None),
+                (Some(1), Some(0)),
+                (Some(2), Some(1)),
+                (Some(3), Some(2)),
+                (None, Some(3)),
+            ]
+        );
+        assert_eq!(
+            pipeline_1f1b_schedule(0, 3, 4).unwrap(),
+            vec![
+                (Some(0), None),
+                (Some(1), None),
+                (Some(2), Some(0)),
+                (Some(3), Some(1)),
+                (None, Some(2)),
+                (None, Some(3)),
+            ]
+        );
+        assert_eq!(
+            pipeline_1f1b_schedule(2, 3, 4).unwrap(),
+            vec![
+                (Some(0), Some(0)),
+                (Some(1), Some(1)),
+                (Some(2), Some(2)),
+                (Some(3), Some(3)),
+            ]
+        );
+    }
+
+    #[test]
+    fn projection_sharded_lora_rank_does_not_require_tp_divisibility() {
+        validate_lora_rank_for_tp(
+            3,
+            2,
+            &[
+                checkpoint::LoraTpShardLayout::ColumnParallel,
+                checkpoint::LoraTpShardLayout::RoutedExpertFusedGateUp,
+            ],
+        )
+        .unwrap();
+        assert!(
+            validate_lora_rank_for_tp(3, 2, &[checkpoint::LoraTpShardLayout::LatentRank],).is_err()
+        );
+    }
+
+    #[test]
+    fn native_optimizer_accepts_gradient_clipping() {
+        validate_native_optimizer_options(Some(1.0)).unwrap();
+        assert!(validate_native_optimizer_options(Some(0.0)).is_err());
+        assert!(validate_native_optimizer_options(Some(f32::NAN)).is_err());
+        validate_native_optimizer_options(None).unwrap();
+    }
+}
+
 // ──────────────────────────────────────────────────────────────────────
 // Core training implementation
 // ──────────────────────────────────────────────────────────────────────
@@ -193,11 +353,17 @@ fn train_impl(
     run_paths: &RunPaths,
     ep_shard: Option<EpShard>,
 ) -> Result<Qwen36LoraSftSummary> {
-    let model_path = config.model.model_path.as_ref()
+    validate_native_optimizer_options(config.train.max_grad_norm)?;
+    let model_path = config
+        .model
+        .model_path
+        .as_ref()
         .ok_or_else(|| anyhow!("model.model_path required"))?;
-    let model_path = resolve_qwen36_model_path(model_path)?;
+    let model_path = std::fs::canonicalize(resolve_qwen36_model_path(model_path)?)
+        .with_context(|| format!("canonicalize Qwen model path {}", model_path.display()))?;
     let runtime_config = read_qwen36_runtime_config(&model_path)?;
     let lora_config = lora_config_from_config(config)?;
+    validate_lora_targets(&runtime_config, &lora_config)?;
     let device = match config.train.device {
         rustrain_core::runtime::Device::Cuda => {
             // EP mode: use LOCAL_RANK to select the correct GPU
@@ -214,56 +380,449 @@ fn train_impl(
         rustrain_core::runtime::DType::Bf16 => Kind::BFloat16,
         rustrain_core::runtime::DType::Fp32 => Kind::Float,
     };
+    if compute_kind != Kind::BFloat16 {
+        bail!(
+            "native Qwen3.5/3.6 LoRA currently supports bf16 only; {:?} would not be updated by the fused Adam kernel",
+            config.train.dtype
+        );
+    }
+    let seed = i64::try_from(config.run.seed)
+        .context("run.seed exceeds libtorch's signed 64-bit seed range")?;
+    tch::manual_seed(seed);
 
     let shard_ref = ep_shard.as_ref();
     let is_ep = shard_ref.is_some();
-    let world_size = shard_ref.map(|s| s.world_size).unwrap_or(1);
-    let rank = shard_ref.map(|s| s.rank).unwrap_or(0);
+    let env_enabled = |name: &str| {
+        std::env::var(name)
+            .map(|value| !value.is_empty() && value != "0")
+            .unwrap_or(false)
+    };
+    let ep_a2a = env_enabled("QWEN36_EP_A2A");
+    let ep_a2a_sharded = env_enabled("QWEN36_EP_A2A_SHARDED");
+    if ep_a2a_sharded && !is_ep {
+        bail!("QWEN36_EP_A2A_SHARDED=1 requires expert-parallel training");
+    }
+    if ep_a2a_sharded && !ep_a2a {
+        bail!("QWEN36_EP_A2A_SHARDED=1 requires QWEN36_EP_A2A=1");
+    }
+    let env_world_size = std::env::var("WORLD_SIZE")
+        .ok()
+        .and_then(|s| s.parse::<usize>().ok())
+        .unwrap_or(1);
+    let env_rank = std::env::var("RANK")
+        .ok()
+        .and_then(|s| s.parse::<usize>().ok())
+        .unwrap_or(0);
+    let world_size = env_world_size;
+    let rank = env_rank;
+    let distributed_generation = if world_size > 1 {
+        let attempt = std::env::var("RUSTRAIN_ATTEMPT_ID")
+            .context("distributed CLI training requires launcher-provided RUSTRAIN_ATTEMPT_ID")?;
+        Some(format!(
+            "{attempt}.adapter-final.{ordinal:020}",
+            ordinal = 0
+        ))
+    } else {
+        None
+    };
+    let tp_size = config.parallel.tensor_model_parallel_size;
+    let pp_size = config.parallel.pipeline_model_parallel_size;
+    let cp_size = config.parallel.context_parallel_size;
+    let env_tp_size = std::env::var("TP_SIZE")
+        .or_else(|_| std::env::var("RUSTRAIN_TP_SIZE"))
+        .or_else(|_| std::env::var("TENSOR_MODEL_PARALLEL_SIZE"))
+        .ok()
+        .and_then(|s| s.parse::<usize>().ok())
+        .unwrap_or(1);
+    if env_tp_size > 1 && env_tp_size != tp_size {
+        bail!(
+            "TP_SIZE environment ({env_tp_size}) does not match config tensor_model_parallel_size ({tp_size})"
+        );
+    }
+    let configured_dp_size = std::env::var("DP_SIZE")
+        .or_else(|_| std::env::var("RUSTRAIN_DP_SIZE"))
+        .or_else(|_| std::env::var("DATA_PARALLEL_SIZE"))
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(config.parallel.data_parallel_size);
+    if is_ep && configured_dp_size != config.parallel.data_parallel_size {
+        bail!(
+            "DP_SIZE environment ({configured_dp_size}) does not match config data_parallel_size ({})",
+            config.parallel.data_parallel_size
+        );
+    }
+    let dense_model_parallel_size = tp_size
+        .checked_mul(pp_size)
+        .and_then(|size| size.checked_mul(cp_size))
+        .context("Qwen model-parallel size overflow")?;
+    let dp_size = if !is_ep && configured_dp_size == 1 && world_size > dense_model_parallel_size {
+        world_size
+            .checked_div(dense_model_parallel_size)
+            .filter(|value| value * dense_model_parallel_size == world_size)
+            .ok_or_else(|| {
+                anyhow!(
+                    "WORLD_SIZE={world_size} is not divisible by TPxPPxCP={dense_model_parallel_size}"
+                )
+            })?
+    } else {
+        configured_dp_size
+    };
+    let rank_order = std::env::var("RUSTRAIN_PARALLEL_ORDER")
+        .or_else(|_| std::env::var("PARALLEL_ORDER"))
+        .unwrap_or_else(|_| DEFAULT_RANK_ORDER.to_string());
+    let cp_full_attention = env_enabled("QWEN36_CP_FULL_ATTENTION_KV_GATHER")
+        || env_enabled("QWEN36_CP_FULL_ATTENTION_RING");
+    let cp_ring_attention = env_enabled("QWEN36_CP_FULL_ATTENTION_RING");
+    if cp_size != 1
+        && ((!cp_ring_attention && cp_size != 2)
+            || tp_size != 1
+            || pp_size != 1
+            || dp_size != 1
+            || is_ep
+            || !cp_full_attention)
+    {
+        bail!(
+            "native Qwen CP requires ring attention for CP>2, otherwise CP2, with TP=EP=DP=PP=1 and one full-attention CP flag (got TP={} EP={} DP={} PP={} CP={})",
+            tp_size,
+            if is_ep { 1 } else { config.parallel.expert_model_parallel_size },
+            dp_size,
+            pp_size,
+            cp_size
+        );
+    }
+    let parallel_topology = if let Some(shard) = shard_ref {
+        Some(ParallelTopology::with_order(
+            tp_size,
+            pp_size,
+            dp_size,
+            shard.world_size,
+            cp_size,
+            &rank_order,
+        )?)
+    } else {
+        if config.parallel.expert_model_parallel_size != 1 {
+            bail!(
+                "native dense Qwen LoRA requires EP=1: TP={} WORLD_SIZE={} PP={} DP={} EP={} CP={}",
+                tp_size,
+                world_size,
+                pp_size,
+                dp_size,
+                config.parallel.expert_model_parallel_size,
+                cp_size
+            );
+        }
+        Some(ParallelTopology::with_order(
+            tp_size,
+            pp_size,
+            dp_size,
+            1,
+            cp_size,
+            &rank_order,
+        )?)
+    };
+    if let Some(topology) = parallel_topology.as_ref() {
+        topology.validate_world_size(world_size)?;
+        topology.coordinates(rank)?;
+    }
+    let sequence_parallel = env_enabled("QWEN36_SEQUENCE_PARALLEL");
+    if sequence_parallel
+        && (runtime_config.is_moe
+            || runtime_config.has_vision
+            || tp_size != 2
+            || pp_size != 1
+            || cp_size != 1
+            || dp_size != 1
+            || runtime_config.mtp_num_hidden_layers > 0
+            || runtime_config.router_aux_loss_coef != 0.0)
+    {
+        bail!(
+            "QWEN36_SEQUENCE_PARALLEL=1 currently requires dense text-only fixed-LoRA TP2 with PP=CP=DP=1, MTP disabled, and router_aux_loss_coef=0"
+        );
+    }
+    let is_data_parallel = dp_size > 1;
+    unsafe {
+        std::env::set_var("TP_SIZE", tp_size.to_string());
+        std::env::set_var(
+            "EP_SIZE",
+            shard_ref
+                .map(|shard| shard.world_size)
+                .unwrap_or(1)
+                .to_string(),
+        );
+        std::env::set_var("DP_SIZE", dp_size.to_string());
+        std::env::set_var("CP_SIZE", cp_size.to_string());
+        std::env::set_var("PP_SIZE", pp_size.to_string());
+        std::env::set_var(
+            "RUSTRAIN_DATA_PARALLEL",
+            if is_data_parallel { "1" } else { "0" },
+        );
+    }
+    let tp_rank = parallel_topology
+        .as_ref()
+        .map(|topology| topology.tensor_rank(rank))
+        .transpose()?
+        .unwrap_or(0);
+    let dp_rank = parallel_topology
+        .as_ref()
+        .map(|topology| topology.data_rank(rank))
+        .transpose()?
+        .unwrap_or(0);
+    let cp_rank = parallel_topology
+        .as_ref()
+        .map(|topology| topology.context_rank(rank))
+        .transpose()?
+        .unwrap_or(0);
+    let pp_rank = parallel_topology
+        .as_ref()
+        .map(|topology| topology.pipeline_rank(rank))
+        .transpose()?
+        .unwrap_or(0);
+    let ep_rank = parallel_topology
+        .as_ref()
+        .map(|topology| topology.expert_rank(rank))
+        .transpose()?
+        .unwrap_or(0);
+    let ep_size = parallel_topology
+        .as_ref()
+        .map(ParallelTopology::expert_model_parallel_size)
+        .unwrap_or(1);
+    let stage = PipelineStageLayout::new(runtime_config.num_hidden_layers, pp_rank, pp_size)?;
+    if pp_size > 1 && runtime_config.has_vision {
+        bail!("pipeline-parallel Qwen training does not yet support the vision encoder");
+    }
+    if pp_size > 1 && runtime_config.mtp_num_hidden_layers > 0 {
+        bail!("pipeline-parallel Qwen training does not yet support MTP layers");
+    }
+    let global_native_slots = native_lora_slots(&runtime_config, &lora_config);
+    let native_slots = stage_lora_slots(&global_native_slots, &stage);
+    let active_layouts = native_slots
+        .iter()
+        .filter(|slot| slot.active)
+        .map(|slot| checkpoint::lora_tp_shard_layout(slot.module, &runtime_config))
+        .collect::<Vec<_>>();
+    validate_lora_rank_for_tp(lora_config.rank, tp_size, &active_layouts)?;
+    let is_expert_parallel = ep_size > 1;
+    unsafe {
+        std::env::set_var("RUSTRAIN_TP_RANK", tp_rank.to_string());
+        std::env::set_var("RUSTRAIN_CP_RANK", cp_rank.to_string());
+        std::env::set_var("RUSTRAIN_EP_RANK", ep_rank.to_string());
+        std::env::set_var("RUSTRAIN_DP_RANK", dp_rank.to_string());
+        std::env::set_var("RUSTRAIN_PP_RANK", pp_rank.to_string());
+    }
+    if is_data_parallel || is_expert_parallel || tp_size > 1 {
+        crate::kernel::CppTrainingContext::set_cuda_device(
+            std::env::var("LOCAL_RANK")
+                .ok()
+                .and_then(|s| s.parse::<i32>().ok())
+                .unwrap_or(0),
+        );
+    }
 
-    // Build needed weight set
-    let needed = build_needed_weights(&runtime_config, &lora_config, shard_ref);
+    // Full attention and GDN follow head-aligned ColumnParallel input
+    // projections and RowParallel output projections. Dense MLP additionally
+    // shards gate/up rows and down columns. Embeddings and the LM head shard
+    // contiguous vocabulary rows and remain tied through the C++ context when
+    // the model uses tied word embeddings.
+    let base_tp_attention = tp_size > 1;
+    let base_tp_mlp = tp_size > 1;
+    // MTP loss still uses a replicated vocabulary projection. Its prediction
+    // layers are tensor parallel, while embedding/LM-head sharding remains a
+    // separate capability boundary.
+    let vocab_parallel = tp_size > 1 && runtime_config.mtp_num_hidden_layers == 0;
+    if vocab_parallel
+        && (runtime_config.vocab_size <= 0 || runtime_config.vocab_size % tp_size as i64 != 0)
+    {
+        bail!(
+            "vocab_size={} must be divisible by TP_SIZE={tp_size} for vocabulary parallelism",
+            runtime_config.vocab_size
+        );
+    }
+    if base_tp_attention {
+        if runtime_config.num_attention_heads <= 0
+            || runtime_config.num_attention_heads % tp_size as i64 != 0
+            || runtime_config.num_key_value_heads <= 0
+            || runtime_config.num_key_value_heads % tp_size as i64 != 0
+            || runtime_config.num_attention_heads % runtime_config.num_key_value_heads != 0
+        {
+            bail!(
+                "full-attention heads (q={}, kv={}) must preserve GQA groups and be divisible by TP_SIZE={tp_size}",
+                runtime_config.num_attention_heads,
+                runtime_config.num_key_value_heads
+            );
+        }
+        let rotary_dim =
+            (runtime_config.head_dim as f64 * runtime_config.partial_rotary_factor) as i64;
+        if runtime_config.head_dim <= 0
+            || rotary_dim < 0
+            || rotary_dim > runtime_config.head_dim
+            || rotary_dim % 2 != 0
+        {
+            bail!(
+                "full-attention head_dim={} and partial_rotary_factor={} produce invalid rotary_dim={rotary_dim}",
+                runtime_config.head_dim,
+                runtime_config.partial_rotary_factor
+            );
+        }
+        if runtime_config
+            .layer_types
+            .iter()
+            .any(|layer| *layer == LayerType::LinearAttention)
+        {
+            if runtime_config.linear_num_key_heads <= 0
+                || runtime_config.linear_num_value_heads <= 0
+                || runtime_config.linear_num_value_heads % runtime_config.linear_num_key_heads != 0
+                || runtime_config.linear_num_key_heads % tp_size as i64 != 0
+                || runtime_config.linear_num_value_heads % tp_size as i64 != 0
+                || runtime_config.linear_key_head_dim != 128
+                || runtime_config.linear_value_head_dim != 128
+                || runtime_config.linear_conv_kernel_dim <= 0
+            {
+                bail!(
+                    "linear-attention TP requires k/v heads divisible by TP_SIZE with preserved value-head groups, 128-wide key/value heads, and a positive conv kernel: k_heads={}, v_heads={}, key_dim={}, value_dim={}, conv_kernel={}, tp={tp_size}",
+                    runtime_config.linear_num_key_heads,
+                    runtime_config.linear_num_value_heads,
+                    runtime_config.linear_key_head_dim,
+                    runtime_config.linear_value_head_dim,
+                    runtime_config.linear_conv_kernel_dim,
+                );
+            }
+        }
+    }
+    if base_tp_mlp {
+        let intermediates = if runtime_config.is_moe {
+            vec![
+                ("routed expert", runtime_config.moe_intermediate_size),
+                (
+                    "shared expert",
+                    runtime_config.shared_expert_intermediate_size,
+                ),
+            ]
+        } else {
+            vec![("dense", runtime_config.intermediate_size)]
+        };
+        for (kind, intermediate) in intermediates {
+            if intermediate <= 0 || intermediate % tp_size as i64 != 0 {
+                bail!(
+                    "{kind} intermediate_size={intermediate} must be divisible by TP_SIZE={tp_size}"
+                );
+            }
+        }
+    }
+    if runtime_config.mtp_num_hidden_layers > 0 && ep_size > 1 {
+        bail!(
+            "fixed-LoRA MTP with EP_SIZE={ep_size} remains fail-closed until its MTP/main denominator uses global source token counts; dynamic multi-LoRA EP MTP is supported through the native API"
+        );
+    }
+
+    // Load only the frozen weights owned by this physical pipeline stage.
+    let needed = stage_needed_weights(&runtime_config, &stage);
 
     // Stagger loading for EP to avoid OOM
     if is_ep {
         std::thread::sleep(std::time::Duration::from_secs(rank as u64 * 5));
     }
 
-    info!("loading {} weight tensors from {}", needed.len(), model_path.display());
+    info!(
+        "loading {} weight tensors from {}",
+        needed.len(),
+        model_path.display()
+    );
     let weights = read_safetensors_dir_filtered(&model_path, &needed)?;
 
-    // Move to device — for EP, narrow expert tensors on CPU first to save GPU memory
+    // Apply orthogonal EP and TP shards on CPU before moving weights to CUDA.
     let mut weights_gpu = BTreeMap::new();
-    if let Some(shard) = shard_ref {
-        // EP mode: narrow expert tensors on CPU before transferring to GPU
-        let num_experts = runtime_config.num_experts as i64;
-        for (name, tensor) in &weights {
-            // Check if this is an expert tensor that needs narrowing
-            let needs_narrow = name.contains(".mlp.experts.gate_up_proj")
-                || name.contains(".mlp.experts.down_proj");
-            if needs_narrow && tensor.size()[0] == num_experts {
-                // Narrow to local shard: [rank * experts_per_rank, ...]
-                let narrowed = tensor
+    let num_experts = runtime_config.num_experts as i64;
+    for (name, tensor) in &weights {
+        let needs_expert_narrow = shard_ref.is_some()
+            && (name.contains(".mlp.experts.gate_up_proj")
+                || name.contains(".mlp.experts.down_proj"));
+        let expert_shard = if needs_expert_narrow && tensor.size()[0] == num_experts {
+            let shard = shard_ref.expect("expert shard checked above");
+            Some(
+                tensor
                     .narrow(0, shard.expert_start as i64, shard.experts_per_rank as i64)
-                    .contiguous()
-                    .to_device(device)
-                    .to_kind(compute_kind);
-                weights_gpu.insert(name.clone(), narrowed);
-            } else if needs_narrow && tensor.size()[0] != num_experts {
-                // Already narrowed (e.g., MTP layer experts) — load as-is
-                weights_gpu.insert(name.clone(), tensor.to_device(device).to_kind(compute_kind));
+                    .contiguous(),
+            )
+        } else {
+            None
+        };
+        let expert_or_full = expert_shard.as_ref().unwrap_or(tensor);
+        let moe_tp_shard = if runtime_config.is_moe && base_tp_mlp {
+            crate::kernel::shard_moe_mlp_weight_for_tp(name, expert_or_full, tp_size, tp_rank)?
+        } else {
+            None
+        };
+        let vocab_shard = if expert_shard.is_none() && moe_tp_shard.is_none() && vocab_parallel {
+            crate::kernel::shard_vocab_weight_for_tp(
+                name,
+                tensor,
+                runtime_config.vocab_size,
+                tp_size,
+                tp_rank,
+            )?
+        } else {
+            None
+        };
+        let local_shard = if moe_tp_shard.is_some() {
+            moe_tp_shard
+        } else if expert_shard.is_some() {
+            expert_shard
+        } else if vocab_shard.is_some() {
+            vocab_shard
+        } else if base_tp_attention {
+            let full_attention_shard =
+                crate::kernel::shard_full_attention_weight_for_tp(name, tensor, tp_size, tp_rank)?;
+            let attention_shard = if full_attention_shard.is_some() {
+                full_attention_shard
             } else {
-                // Non-expert weights: replicate on all GPUs
-                weights_gpu.insert(name.clone(), tensor.to_device(device).to_kind(compute_kind));
+                crate::kernel::shard_linear_attention_weight_for_tp(
+                    name,
+                    tensor,
+                    tp_size,
+                    tp_rank,
+                    runtime_config.linear_num_key_heads,
+                    runtime_config.linear_key_head_dim,
+                    runtime_config.linear_num_value_heads,
+                    runtime_config.linear_value_head_dim,
+                )?
+            };
+            if attention_shard.is_some() || !base_tp_mlp {
+                attention_shard
+            } else {
+                crate::kernel::shard_dense_mlp_weight_for_tp(name, tensor, tp_size, tp_rank)?
             }
-        }
-        info!("EP{}: narrowed expert tensors to {} experts per rank", world_size, shard.experts_per_rank);
-    } else {
-        for (name, tensor) in &weights {
-            weights_gpu.insert(name.clone(), tensor.to_device(device).to_kind(compute_kind));
-        }
+        } else {
+            None
+        };
+        let gpu_tensor = local_shard
+            .as_ref()
+            .unwrap_or(tensor)
+            .to_device(device)
+            .to_kind(compute_kind);
+        weights_gpu.insert(name.clone(), gpu_tensor);
+    }
+    if let Some(shard) = shard_ref {
+        info!(
+            ep_size,
+            ep_rank,
+            experts_per_rank = shard.experts_per_rank,
+            "frozen expert weights sharded across EP group"
+        );
+    }
+    if base_tp_attention {
+        info!(
+            tp_size,
+            tp_rank,
+            base_tp_mlp,
+            vocab_parallel,
+            "frozen base TP enabled: attention/GDN, vocabulary, and dense/expert MLP shards"
+        );
     }
 
-    info!("LoRA config: rank={}, alpha={}", lora_config.rank, lora_config.alpha);
+    info!(
+        "LoRA config: rank={}, alpha={}",
+        lora_config.rank, lora_config.alpha
+    );
 
     // Load SFT data
     let tokenizer_path = model_path.join("tokenizer.json");
@@ -281,7 +840,47 @@ fn train_impl(
         let mtp_tensors = read_safetensors_dir_filtered(&model_path, &mtp_needed)?;
         let mut mtp_gpu = BTreeMap::new();
         for (name, tensor) in &mtp_tensors {
-            mtp_gpu.insert(name.clone(), tensor.to_device(device).to_kind(compute_kind));
+            let needs_expert_narrow = shard_ref.is_some()
+                && (name.contains(".mlp.experts.gate_up_proj")
+                    || name.contains(".mlp.experts.down_proj"));
+            let expert_shard = if needs_expert_narrow && tensor.size()[0] == num_experts {
+                let shard = shard_ref.expect("expert shard checked above");
+                Some(
+                    tensor
+                        .narrow(0, shard.expert_start as i64, shard.experts_per_rank as i64)
+                        .contiguous(),
+                )
+            } else {
+                None
+            };
+            let expert_or_full = expert_shard.as_ref().unwrap_or(tensor);
+            let moe_shard = if base_tp_mlp {
+                crate::kernel::shard_moe_mlp_weight_for_tp(
+                    name,
+                    expert_or_full,
+                    tp_size,
+                    tp_rank,
+                )?
+            } else {
+                None
+            };
+            let attention_shard = if moe_shard.is_none() && base_tp_attention {
+                crate::kernel::shard_full_attention_weight_for_tp(name, tensor, tp_size, tp_rank)?
+            } else {
+                None
+            };
+            let dense_shard = if moe_shard.is_none() && attention_shard.is_none() && base_tp_mlp {
+                crate::kernel::shard_dense_mlp_weight_for_tp(name, tensor, tp_size, tp_rank)?
+            } else {
+                None
+            };
+            let local = moe_shard
+                .as_ref()
+                .or(attention_shard.as_ref())
+                .or(dense_shard.as_ref())
+                .or(expert_shard.as_ref())
+                .unwrap_or(tensor);
+            mtp_gpu.insert(name.clone(), local.to_device(device).to_kind(compute_kind));
         }
         for (name, tensor) in mtp_gpu {
             weights_gpu.insert(name, tensor);
@@ -290,35 +889,95 @@ fn train_impl(
 
     let batch_size = config.train.micro_batch_size;
     let max_steps = config.train.max_steps as usize;
+    let gradient_accumulation_steps = config.train.gradient_accumulation_steps;
 
     // ── C++ all-in-C++ training path (required) ──
     // LoRA A/B, Adam optimizer, forward, loss, backward all in C++.
     if !crate::kernel::kernels_available() {
-        bail!("C++ kernels (libqwen36_kernels.so) not found — required for training. Ensure the .so is in LD_LIBRARY_PATH.");
+        bail!(
+            "C++ kernels (libqwen36_kernels.so) not found — required for training. Ensure the .so is in LD_LIBRARY_PATH."
+        );
     }
 
-    let ctx = crate::kernel::CppTrainingContext::new(
-        &weights_gpu, &runtime_config, compute_kind,
+    let expert_start = shard_ref.map(|s| s.expert_start).unwrap_or(0);
+    let expert_count = shard_ref
+        .map(|s| s.experts_per_rank)
+        .unwrap_or(runtime_config.num_experts);
+    let ctx = crate::kernel::CppTrainingContext::new_for_stage(
+        &weights_gpu,
+        &runtime_config,
+        &stage,
+        compute_kind,
         config.train.learning_rate as f64,
         config.train.adam_beta1 as f64,
         config.train.adam_beta2 as f64,
         config.train.adam_eps as f64,
-        lora_config.alpha as f64 / lora_config.rank as f64,  // lora scaling = alpha / rank
+        lora_config.alpha as f64 / lora_config.rank as f64, // lora scaling = alpha / rank
         lora_config.rank as i64,
+        base_tp_attention,
+        base_tp_mlp,
+        vocab_parallel,
+        is_data_parallel,
+        is_expert_parallel,
         &lora_config.target_layers,
-        shard_ref.map(|s| s.expert_start).unwrap_or(0),
-        shard_ref.map(|s| s.experts_per_rank).unwrap_or(0),
+        &lora_config.target_modules,
+        expert_start,
+        expert_count,
     )?;
+    ctx.set_pad_token_id(data.pad_token_id())?;
+    ctx.set_max_grad_norm(config.train.max_grad_norm.map(f64::from).unwrap_or(0.0))?;
+
+    if world_size > 1 {
+        if let Some(topology) = parallel_topology.as_ref() {
+            let tp_color = *topology
+                .tensor_group(rank)?
+                .iter()
+                .min()
+                .ok_or_else(|| anyhow!("empty TP process group"))?;
+            let ep_color = *topology
+                .expert_group(rank)?
+                .iter()
+                .min()
+                .ok_or_else(|| anyhow!("empty EP process group"))?;
+            let cp_color = *topology
+                .context_group(rank)?
+                .iter()
+                .min()
+                .ok_or_else(|| anyhow!("empty CP process group"))?;
+            let dp_color = *topology
+                .data_group(rank)?
+                .iter()
+                .min()
+                .ok_or_else(|| anyhow!("empty DP process group"))?;
+            let pp_color = *topology
+                .pipeline_group(rank)?
+                .iter()
+                .min()
+                .ok_or_else(|| anyhow!("empty PP process group"))?;
+            ctx.init_parallel_nccl(
+                rank, world_size, tp_rank, tp_size, tp_color, cp_rank, cp_size, cp_color, ep_rank,
+                ep_size, ep_color, dp_rank, dp_size, dp_color, pp_rank, pp_size, pp_color,
+            )?;
+        } else {
+            let ret = ctx.init_nccl();
+            if ret != 0 {
+                bail!("C++ NCCL init failed (code {})", ret);
+            }
+        }
+        info!(
+            rank,
+            world_size, is_ep, "NCCL communicator created for Qwen LoRA parallel training"
+        );
+    }
     info!("C++ TrainingContext: {} LoRA params", ctx.lora_count());
 
     // Set MTP weights if available
     if runtime_config.mtp_num_hidden_layers > 0 {
-        ctx.set_mtp_weights(
-            &weights_gpu, &runtime_config,
-            shard_ref.map(|s| s.expert_start).unwrap_or(0),
-            shard_ref.map(|s| s.experts_per_rank).unwrap_or(0),
-        )?;
-        info!("C++ TrainingContext: MTP weights set ({} layers)", runtime_config.mtp_num_hidden_layers);
+        ctx.set_mtp_weights(&weights_gpu, &runtime_config, expert_start, expert_count)?;
+        info!(
+            "C++ TrainingContext: MTP weights set ({} layers)",
+            runtime_config.mtp_num_hidden_layers
+        );
     }
 
     // Enable gradient checkpointing if env var set
@@ -328,74 +987,369 @@ fn train_impl(
         info!("C++ TrainingContext: gradient checkpointing ON (group_size={group_size})");
     }
 
+    let native_count = ctx.lora_count() as usize;
+    if native_slots.len() != native_count {
+        bail!(
+            "pipeline-local LoRA registry count {} does not match native slot count {native_count}",
+            native_slots.len()
+        );
+    }
+    let mut start_step = 0usize;
+    let mut resumed_loss = 0.0_f64;
+    if let Some(resume_path) = config.train.resume_from.as_ref() {
+        let parallel = checkpoint::ParallelCheckpointManifest::from_topology(
+            world_size,
+            rank,
+            parallel_topology
+                .as_ref()
+                .context("Qwen checkpoint restore is missing parallel topology")?,
+        )?;
+        let data = checkpoint::load_checkpoint_for_topology(resume_path, &parallel)?;
+        let checkpoint_model_path =
+            std::fs::canonicalize(&data.manifest.model_path).with_context(|| {
+                format!(
+                    "canonicalize checkpoint base model path {}",
+                    data.manifest.model_path
+                )
+            })?;
+        if checkpoint_model_path != model_path {
+            bail!(
+                "checkpoint base model {} does not match configured model {}",
+                checkpoint_model_path.display(),
+                model_path.display()
+            );
+        }
+        if !data.dynamic_adapters.is_empty() {
+            bail!("Qwen CLI fixed-LoRA training cannot resume dynamic adapter state");
+        }
+        if data.manifest.lora_rank != lora_config.rank
+            || (data.manifest.lora_alpha - lora_config.alpha).abs() > 1e-12
+        {
+            bail!(
+                "checkpoint LoRA rank/alpha {}/{} does not match config {}/{}",
+                data.manifest.lora_rank,
+                data.manifest.lora_alpha,
+                lora_config.rank,
+                lora_config.alpha
+            );
+        }
+        let expected_identities = native_slots
+            .iter()
+            .filter(|slot| slot.active)
+            .map(|slot| checkpoint::LoraSlotIdentity {
+                index: slot.global_index,
+                layer: slot.layer,
+                module: slot.module.cpp_name().to_string(),
+            })
+            .collect::<Vec<_>>();
+        if world_size > 1 {
+            checkpoint::validate_fixed_tp_resume(
+                &data.manifest,
+                &active_layouts,
+                &expected_identities,
+            )?;
+        }
+        let active_slot_indices = native_slots
+            .iter()
+            .filter(|slot| slot.active)
+            .map(|slot| slot.local_index)
+            .collect::<Vec<_>>();
+        let restore_slot_indices = checkpoint::fixed_restore_slot_indices(
+            data.lora_a.len(),
+            data.lora_b.len(),
+            &active_slot_indices,
+            native_count,
+        )?;
+        for ((a, b), &slot_index) in data
+            .lora_a
+            .iter()
+            .zip(&data.lora_b)
+            .zip(&restore_slot_indices)
+        {
+            ctx.set_lora_tensor(slot_index as i64, false, a)?;
+            ctx.set_lora_tensor(slot_index as i64, true, b)?;
+        }
+        if data.adam_m.is_empty() != data.adam_v.is_empty() {
+            bail!(
+                "checkpoint fixed optimizer m/v count mismatch: {}/{}",
+                data.adam_m.len(),
+                data.adam_v.len()
+            );
+        }
+        if data.adam_m.is_empty() && data.manifest.effective_fixed_optimizer_step() > 0 {
+            bail!(
+                "checkpoint fixed optimizer step {} has no Adam state",
+                data.manifest.effective_fixed_optimizer_step()
+            );
+        }
+        if !data.adam_m.is_empty() {
+            let expected_saved_optimizer_count = restore_slot_indices.len().saturating_mul(2);
+            if data.adam_m.len() != expected_saved_optimizer_count
+                || data.adam_v.len() != expected_saved_optimizer_count
+            {
+                bail!(
+                    "checkpoint fixed optimizer count mismatch: m={}, v={}, expected={expected_saved_optimizer_count}",
+                    data.adam_m.len(),
+                    data.adam_v.len()
+                );
+            }
+            let (mut all_adam_m, mut all_adam_v) = ctx.export_optimizer_state()?;
+            let expected_native_optimizer_count = native_count.saturating_mul(2);
+            if all_adam_m.len() != expected_native_optimizer_count
+                || all_adam_v.len() != expected_native_optimizer_count
+            {
+                bail!(
+                    "native fixed optimizer count mismatch: m={}, v={}, expected={expected_native_optimizer_count}",
+                    all_adam_m.len(),
+                    all_adam_v.len()
+                );
+            }
+            for (saved_slot, &native_slot) in restore_slot_indices.iter().enumerate() {
+                let saved = saved_slot.saturating_mul(2);
+                let native = native_slot.saturating_mul(2);
+                all_adam_m[native] = data.adam_m[saved].shallow_clone();
+                all_adam_m[native + 1] = data.adam_m[saved + 1].shallow_clone();
+                all_adam_v[native] = data.adam_v[saved].shallow_clone();
+                all_adam_v[native + 1] = data.adam_v[saved + 1].shallow_clone();
+            }
+            let imported = ctx.import_optimizer_state(&all_adam_m, &all_adam_v)?;
+            if imported != expected_native_optimizer_count as i64 {
+                bail!(
+                    "native fixed optimizer import restored {imported} tensors, expected {expected_native_optimizer_count}"
+                );
+            }
+        }
+        let native_step = i64::try_from(data.manifest.effective_fixed_optimizer_step())
+            .context("checkpoint step exceeds the native optimizer range")?;
+        ctx.set_step_count(native_step)?;
+        start_step = usize::try_from(data.manifest.step)
+            .context("checkpoint step exceeds the CLI training range")?;
+        if start_step > max_steps {
+            bail!("checkpoint step {start_step} exceeds configured max_steps {max_steps}");
+        }
+        resumed_loss = data.manifest.loss;
+        info!(
+            step = start_step,
+            loss = resumed_loss,
+            path = %resume_path.display(),
+            "Qwen fixed LoRA checkpoint restored"
+        );
+    }
+
     // Training loop
-    let mut initial_loss = 0.0_f64;
-    let mut final_loss = 0.0_f64;
-
-    for step in 0..max_steps {
-        let data_start = (step * batch_size) % data.len();
-        let sft_batch = data.batch(data_start, batch_size);
-        let (input_ids, target_mask) = sft_batch.to_tensors(device, compute_kind);
-
-        // Build attention mask: 1 for real tokens, 0 for padding
-        // target_mask > 0 means the token is either response (loss) or prompt (no loss but attend)
-        // We need: 1 for all non-padding tokens, 0 for padding tokens
-        // The SFT batch's target_mask is: 0=prompt, 1=response, but padding has mask=0 too.
-        // We need attention_mask = (target_mask >= 0).to(float) but that's all 1s.
-        // Actually: padding tokens have target_mask=0 AND are after EOS.
-        // The SftBatch already has padding at the end with mask=0.
-        // We need: attention_mask = 1 where token is NOT padding.
-        // Since prompt tokens have mask=0 and response tokens have mask=1,
-        // but padding also has mask=0, we can't distinguish prompt from padding using mask alone.
-        // Solution: use the pad_token_id to build attention mask from input_ids.
-        let pad_id = data.pad_token_id();
-        let attention_mask = input_ids.ne(pad_id).to_kind(Kind::Float).unsqueeze(0);  // [1, seq]
-
-        // C++ all-in-C++ path: single call does forward + loss + backward + Adam
-        let loss_value = ctx.train_step(&input_ids, &target_mask, &attention_mask)?;
-        if step == 0 { initial_loss = loss_value; }
+    let mut initial_loss = resumed_loss;
+    let mut final_loss = resumed_loss;
+    let mut observed_loss = false;
+
+    for step in start_step..max_steps {
+        let load_microbatch = |accumulation_index: usize| {
+            let micro_step = step * gradient_accumulation_steps + accumulation_index;
+            let (source_rank, source_count) =
+                training_source_coordinate(dp_rank, dp_size, ep_rank, ep_size, ep_a2a_sharded);
+            let data_start =
+                (micro_step * batch_size * source_count + source_rank * batch_size) % data.len();
+            let sft_batch = data.batch(data_start, batch_size);
+            let (input_ids, target_mask) = sft_batch.to_tensors(device, compute_kind);
+
+            // The native C++ context derives the padding mask from input IDs,
+            // keeping GPU mask construction out of this Rust hot loop.
+            (input_ids, target_mask)
+        };
+
+        let loss_value = if pp_size > 1 {
+            let window_id = i64::try_from(step).context("pipeline window id exceeds i64")?;
+            let num_microbatches = i64::try_from(gradient_accumulation_steps)
+                .context("pipeline microbatch count exceeds i64")?;
+            ctx.pipeline_begin_v1(window_id, num_microbatches)?;
+            let schedule_result = (|| -> Result<f64> {
+                let gradient_scale = 1.0 / gradient_accumulation_steps as f64;
+                for (forward_mb, backward_mb) in
+                    pipeline_1f1b_schedule(pp_rank, pp_size, gradient_accumulation_steps)?
+                {
+                    if let Some(microbatch) = forward_mb {
+                        let (input_ids, target_mask) =
+                            load_microbatch(microbatch as usize);
+                        ctx.pipeline_tick_v1(
+                            window_id,
+                            Some(microbatch),
+                            backward_mb,
+                            Some(&input_ids),
+                            Some(&target_mask),
+                            None,
+                            gradient_scale,
+                        )?;
+                    } else {
+                        ctx.pipeline_tick_v1(
+                            window_id,
+                            None,
+                            backward_mb,
+                            None,
+                            None,
+                            None,
+                            gradient_scale,
+                        )?;
+                    }
+                }
+                Ok(ctx.pipeline_finish_v1()?.loss)
+            })();
+            match schedule_result {
+                Ok(loss) => loss,
+                Err(error) => {
+                    let _ = ctx.pipeline_abort_v1();
+                    return Err(error);
+                }
+            }
+        } else {
+            let mut loss = 0.0;
+            for accumulation_index in 0..gradient_accumulation_steps {
+                let (input_ids, target_mask) = load_microbatch(accumulation_index);
+
+                loss += ctx.train_micro_step(
+                    &input_ids,
+                    &target_mask,
+                    None,
+                    1.0 / gradient_accumulation_steps as f64,
+                    accumulation_index + 1 == gradient_accumulation_steps,
+                )? / gradient_accumulation_steps as f64;
+            }
+            loss
+        };
+        if !observed_loss {
+            initial_loss = loss_value;
+            observed_loss = true;
+        }
         final_loss = loss_value;
         if step % 10 == 0 || step == max_steps - 1 {
             info!("step {step}/{max_steps} loss={loss_value:.6}");
         }
     }
 
-    // Save adapter — export LoRA A/B from C++ to safetensors
-    let adapter_path = run_paths.root.join("adapter.safetensors");
-    {
-        let mut named_tensors: BTreeMap<String, Tensor> = BTreeMap::new();
-        for i in 0..ctx.lora_count() {
-            if let (Some(a), Some(b)) = (ctx.get_lora_a(i), ctx.get_lora_b(i)) {
-                named_tensors.insert(format!("lora_a_{i}"), a.to_kind(Kind::Float).to_device(tch::Device::Cpu));
-                named_tensors.insert(format!("lora_b_{i}"), b.to_kind(Kind::Float).to_device(tch::Device::Cpu));
-            }
+    let (adapter_path, trainable_params) = if world_size > 1 {
+        let generation = distributed_generation
+            .as_deref()
+            .expect("distributed save generation was validated before training");
+        let parallel = checkpoint::ParallelCheckpointManifest::from_topology(
+            world_size,
+            rank,
+            parallel_topology
+                .as_ref()
+                .context("distributed CLI export is missing parallel topology")?,
+        )?;
+        let (all_adam_m, all_adam_v) = ctx.export_optimizer_state()?;
+        let expected_optimizer_count = native_count.saturating_mul(2);
+        if all_adam_m.len() != expected_optimizer_count
+            || all_adam_v.len() != expected_optimizer_count
+        {
+            bail!(
+                "fixed optimizer state count mismatch: m={}, v={}, expected={expected_optimizer_count}",
+                all_adam_m.len(),
+                all_adam_v.len()
+            );
         }
-        use std::io::Write;
-        let mut tensors_data = Vec::new();
-        let mut header = serde_json::Map::new();
-        let mut offset = 0u64;
-        for (name, tensor) in &named_tensors {
-            let t = tensor.contiguous().to_kind(Kind::Float);
-            let shape: Vec<i64> = t.size().iter().copied().collect();
-            let data: Vec<f32> = Vec::<f32>::try_from(&t.reshape([-1]))?;
-            let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
-            header.insert(name.clone(), serde_json::json!({"dtype":"F32","shape":shape,"data_offsets":[offset,offset+bytes.len() as u64]}));
-            offset += bytes.len() as u64;
-            tensors_data.push(bytes);
+        let active_slots = native_slots
+            .iter()
+            .filter(|slot| slot.active)
+            .collect::<Vec<_>>();
+        let mut lora_a = Vec::with_capacity(active_slots.len());
+        let mut lora_b = Vec::with_capacity(active_slots.len());
+        let mut adam_m = Vec::with_capacity(active_slots.len().saturating_mul(2));
+        let mut adam_v = Vec::with_capacity(active_slots.len().saturating_mul(2));
+        let mut layouts = Vec::with_capacity(active_slots.len());
+        let mut identities = Vec::with_capacity(active_slots.len());
+        for slot in active_slots {
+            lora_a.push(ctx.get_lora_a(slot.local_index as i64).with_context(|| {
+                format!(
+                    "native LoRA slot {} (global {}) is missing A",
+                    slot.local_index, slot.global_index
+                )
+            })?);
+            lora_b.push(ctx.get_lora_b(slot.local_index as i64).with_context(|| {
+                format!(
+                    "native LoRA slot {} (global {}) is missing B",
+                    slot.local_index, slot.global_index
+                )
+            })?);
+            let optimizer_index = slot.local_index.saturating_mul(2);
+            adam_m.push(all_adam_m[optimizer_index].shallow_clone());
+            adam_m.push(all_adam_m[optimizer_index + 1].shallow_clone());
+            adam_v.push(all_adam_v[optimizer_index].shallow_clone());
+            adam_v.push(all_adam_v[optimizer_index + 1].shallow_clone());
+            layouts.push(checkpoint::lora_tp_shard_layout(
+                slot.module,
+                &runtime_config,
+            ));
+            identities.push(checkpoint::LoraSlotIdentity {
+                index: slot.global_index,
+                layer: slot.layer,
+                module: slot.module.cpp_name().to_string(),
+            });
         }
-        let header_str = serde_json::to_string(&serde_json::Value::Object(header))?;
-        let file = std::fs::File::create(&adapter_path).with_context(|| format!("create {}", adapter_path.display()))?;
-        let mut writer = std::io::BufWriter::new(file);
-        writer.write_all(&(header_str.len() as u64).to_le_bytes())?;
-        writer.write_all(header_str.as_bytes())?;
-        for data in &tensors_data { writer.write_all(data)?; }
-        info!("saved adapter to {}", adapter_path.display());
-    }
+        let adapter_dir = run_paths.root.join("adapter");
+        let step = u64::try_from(ctx.get_step_count())
+            .context("native fixed adapter optimizer step is negative")?;
+        let model_path_string = model_path
+            .to_str()
+            .context("Qwen model path is not UTF-8")?
+            .to_string();
+        let count = checkpoint::export_distributed_adapter_checkpoint(
+            &adapter_dir,
+            generation,
+            &parallel,
+            None,
+            |staging| {
+                checkpoint::save_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+                    staging,
+                    step,
+                    step,
+                    final_loss,
+                    &model_path_string,
+                    lora_config.rank,
+                    lora_config.alpha,
+                    &lora_a,
+                    &lora_b,
+                    &adam_m,
+                    &adam_v,
+                    &[],
+                    &layouts,
+                    &identities,
+                    &parallel,
+                    Some(generation),
+                )
+            },
+        )?;
+        (adapter_dir.join("adapter_model.safetensors"), count)
+    } else {
+        // Export every positional native slot, then let the artifact mapper omit
+        // inactive slots and assign stable projection-aware tensor names.
+        let mut exported = Vec::with_capacity(native_count);
+        for index in 0..ctx.lora_count() {
+            let a = ctx
+                .get_lora_a(index)
+                .with_context(|| format!("native LoRA slot {index} is missing A"))?;
+            let b = ctx
+                .get_lora_b(index)
+                .with_context(|| format!("native LoRA slot {index} is missing B"))?;
+            exported.push((a, b));
+        }
+        let artifact = Qwen36AdapterArtifact::from_native_exports(
+            &config.model.name,
+            &config.model.architecture,
+            Some(&model_path),
+            &runtime_config,
+            &lora_config,
+            exported,
+        )?;
+        let count = artifact.tensors.len();
+        let path = run_paths.root.join("adapter_model.safetensors");
+        artifact.save(&run_paths.root)?;
+        (path, count)
+    };
+    info!("saved adapter to {}", adapter_path.display());
 
     Ok(Qwen36LoraSftSummary {
         adapter_output: adapter_path.to_string_lossy().to_string(),
         initial_loss,
         final_loss,
-        trainable_params: ctx.lora_count() as usize * 2,
+        trainable_params,
     })
 }
diff --git a/crates/rustrain-qwen3-6/tests/integration.rs b/crates/rustrain-qwen3-6/tests/integration.rs
index cb291a6d..865c48fe 100644
--- a/crates/rustrain-qwen3-6/tests/integration.rs
+++ b/crates/rustrain-qwen3-6/tests/integration.rs
@@ -1,11 +1,205 @@
 //! Integration tests for Qwen3.6 config parsing and forward pass
 
-use std::collections::BTreeMap;
-
-use rustrain_qwen3_6::config::{read_qwen36_runtime_config, resolve_qwen36_model_path, LayerType};
-use rustrain_qwen3_6::model::{qwen36_forward_from_ids, Qwen36LayerWeights};
+use rustrain_qwen3_6::config::{LayerType, read_qwen36_runtime_config};
+use rustrain_qwen3_6::lora::{
+    Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule, native_lora_slots,
+    validate_lora_targets,
+};
+use rustrain_qwen3_6::model::Qwen36LayerWeights;
 
 const MODEL_PATH: &str = "/vePFS-Mindverse/share/huggingface/hub/models--Qwen--Qwen3.6-35B-A3B/snapshots/995ad96eacd98c81ed38be0c5b274b04031597b0";
+const QWEN35_MODEL_PATH: &str = "/vePFS-Mindverse/share/huggingface/hub/models--Qwen--Qwen3.5-0.8B/snapshots/2fc06364715b967f1860aea9cf38778875588b17";
+
+#[test]
+fn test_native_lora_target_module_contract() {
+    let names = [
+        ("q_proj", "self_attn.q_proj"),
+        ("k_proj", "self_attn.k_proj"),
+        ("v_proj", "self_attn.v_proj"),
+        ("o_proj", "self_attn.o_proj"),
+        ("in_proj_qkv", "linear_attn.in_proj_qkv"),
+        ("in_proj_z", "linear_attn.in_proj_z"),
+        ("out_proj", "linear_attn.out_proj"),
+    ];
+    for (cpp_name, suffix) in names {
+        let module = Qwen36LoraTargetModule::parse(cpp_name).expect("supported target");
+        assert_eq!(module.cpp_name(), cpp_name);
+        assert_eq!(module.suffix(), suffix);
+    }
+    for name in [
+        "in_proj_a",
+        "in_proj_b",
+        "gate_proj",
+        "up_proj",
+        "down_proj",
+        "shared_gate_proj",
+        "shared_up_proj",
+        "shared_down_proj",
+        "experts_gate_up_proj",
+        "experts_down_proj",
+    ] {
+        assert!(Qwen36LoraTargetModule::parse(name).is_ok(), "{name}");
+    }
+}
+
+#[test]
+fn test_routed_expert_adapter_roundtrip_uses_rank3_local_shards() {
+    let runtime =
+        read_qwen36_runtime_config(std::path::Path::new(MODEL_PATH)).expect("Qwen3.6 config parse");
+    let lora = Qwen36LoraConfig {
+        rank: 2,
+        alpha: 8.0,
+        target_layers: vec![0],
+        target_modules: vec![
+            Qwen36LoraTargetModule::ExpertsGateUpProj,
+            Qwen36LoraTargetModule::ExpertsDownProj,
+        ],
+    };
+    validate_lora_targets(&runtime, &lora).expect("MoE routed expert targets");
+    let slots = native_lora_slots(&runtime, &lora);
+    let exported = slots
+        .iter()
+        .map(|slot| match slot.module {
+            Qwen36LoraTargetModule::ExpertsGateUpProj if slot.active => (
+                tch::Tensor::ones([2, 2, 3], (tch::Kind::Float, tch::Device::Cpu)),
+                tch::Tensor::zeros([2, 4, 2], (tch::Kind::Float, tch::Device::Cpu)),
+            ),
+            Qwen36LoraTargetModule::ExpertsDownProj if slot.active => (
+                tch::Tensor::ones([2, 2, 4], (tch::Kind::Float, tch::Device::Cpu)),
+                tch::Tensor::zeros([2, 3, 2], (tch::Kind::Float, tch::Device::Cpu)),
+            ),
+            _ => (
+                tch::Tensor::zeros([], (tch::Kind::Float, tch::Device::Cpu)),
+                tch::Tensor::zeros([], (tch::Kind::Float, tch::Device::Cpu)),
+            ),
+        })
+        .collect();
+    let artifact = Qwen36AdapterArtifact::from_native_exports(
+        "Qwen3.6-35B-A3B",
+        "qwen3_6_lora_sft",
+        Some(std::path::Path::new(MODEL_PATH)),
+        &runtime,
+        &lora,
+        exported,
+    )
+    .expect("build expert adapter artifact");
+
+    assert_eq!(artifact.tensors.len(), 4);
+    let key = format!(
+        "base_model.model.{}layers.0.mlp.experts.gate_up_proj.lora_A.weight",
+        runtime.weight_prefix
+    );
+    assert_eq!(artifact.tensors[&key].size(), [2, 2, 3]);
+
+    let temp = tempfile::tempdir().expect("temporary expert adapter directory");
+    artifact.save(temp.path()).expect("save expert adapter");
+    let loaded = Qwen36AdapterArtifact::load(temp.path()).expect("reload expert adapter");
+    assert_eq!(loaded.tensors[&key].size(), [2, 2, 3]);
+
+    let dense_runtime = read_qwen36_runtime_config(std::path::Path::new(QWEN35_MODEL_PATH))
+        .expect("Qwen3.5 dense config parse");
+    assert!(validate_lora_targets(&dense_runtime, &lora).is_err());
+}
+
+#[test]
+fn test_adapter_artifact_roundtrip_uses_projection_names() {
+    let runtime = read_qwen36_runtime_config(std::path::Path::new(QWEN35_MODEL_PATH))
+        .expect("Qwen3.5 config parse");
+    let lora = Qwen36LoraConfig {
+        rank: 2,
+        alpha: 4.0,
+        target_layers: vec![0, 3],
+        target_modules: vec![
+            Qwen36LoraTargetModule::InProjQkv,
+            Qwen36LoraTargetModule::QProj,
+        ],
+    };
+    let slots = native_lora_slots(&runtime, &lora);
+    let exported = slots
+        .iter()
+        .map(|slot| {
+            let value = slot.index as f64 + 1.0;
+            (
+                tch::Tensor::full([2, 3], value, (tch::Kind::Float, tch::Device::Cpu)),
+                tch::Tensor::full([5, 2], -value, (tch::Kind::Float, tch::Device::Cpu)),
+            )
+        })
+        .collect();
+    let artifact = Qwen36AdapterArtifact::from_native_exports(
+        "Qwen3.5-0.8B",
+        "qwen3_5_lora_sft",
+        Some(std::path::Path::new(QWEN35_MODEL_PATH)),
+        &runtime,
+        &lora,
+        exported,
+    )
+    .expect("build adapter artifact");
+
+    assert_eq!(artifact.tensors.len(), 4);
+    let linear_key = format!(
+        "base_model.model.{}layers.0.linear_attn.in_proj_qkv.lora_A.weight",
+        runtime.weight_prefix
+    );
+    let full_key = format!(
+        "base_model.model.{}layers.3.self_attn.q_proj.lora_B.weight",
+        runtime.weight_prefix
+    );
+    assert!(artifact.tensors.contains_key(&linear_key));
+    assert!(artifact.tensors.contains_key(&full_key));
+
+    let temp = tempfile::tempdir().expect("temporary adapter directory");
+    let tensor_path = artifact.save(temp.path()).expect("save adapter artifact");
+    assert_eq!(
+        tensor_path.file_name().and_then(|name| name.to_str()),
+        Some("adapter_model.safetensors")
+    );
+    assert!(temp.path().join("rustrain_adapter.json").is_file());
+    let peft_config: serde_json::Value = serde_json::from_slice(
+        &std::fs::read(temp.path().join("adapter_config.json")).expect("read PEFT config"),
+    )
+    .expect("parse PEFT config");
+    assert_eq!(peft_config["peft_type"], "LORA");
+    assert_eq!(peft_config["r"], 2);
+    assert_eq!(
+        peft_config["layers_to_transform"],
+        serde_json::json!([0, 3])
+    );
+    assert_eq!(
+        peft_config["target_modules"],
+        serde_json::json!(["in_proj_qkv", "q_proj"])
+    );
+    let loaded = Qwen36AdapterArtifact::load(temp.path()).expect("reload adapter artifact");
+    assert_eq!(loaded.config, artifact.config);
+    assert_eq!(loaded.tensors.len(), artifact.tensors.len());
+    for (name, expected) in &artifact.tensors {
+        let actual = loaded.tensors.get(name).expect("roundtrip tensor name");
+        assert_eq!(actual.size(), expected.size());
+        assert_eq!(
+            Vec::<f32>::try_from(&actual.reshape([-1])).expect("actual values"),
+            Vec::<f32>::try_from(&expected.reshape([-1])).expect("expected values")
+        );
+    }
+    // A PEFT-only export is accepted when rustrain_adapter.json is absent.
+    std::fs::remove_file(temp.path().join("rustrain_adapter.json"))
+        .expect("remove private rustrain metadata");
+    let peft_only = Qwen36AdapterArtifact::load(temp.path()).expect("load PEFT-only artifact");
+    assert_eq!(peft_only.config.r, 2);
+    assert_eq!(peft_only.config.target_layers, vec![0, 3]);
+}
+
+#[test]
+fn test_qwen35_text_config_parsing() {
+    let config = read_qwen36_runtime_config(std::path::Path::new(QWEN35_MODEL_PATH))
+        .expect("Qwen3.5 config parse");
+    assert_eq!(config.num_hidden_layers, 24);
+    assert_eq!(config.hidden_size, 1024);
+    assert_eq!(config.layer_types[3], LayerType::FullAttention);
+    assert_eq!(config.full_attention_interval, 4);
+    assert_eq!(config.linear_num_key_heads, 16);
+    assert_eq!(config.linear_num_value_heads, 16);
+    assert!(config.attn_output_gate);
+    assert!(config.tie_word_embeddings);
+}
 
 #[test]
 fn test_config_parsing() {
@@ -99,21 +293,18 @@ fn test_forward_single_layer() {
     .map(|suffix| format!("{layer_prefix}.{suffix}"))
     .collect();
 
-    let weights = read_safetensors_dir_filtered(model_path, &needed)
-        .expect("load layer 0 weights");
+    let weights = read_safetensors_dir_filtered(model_path, &needed).expect("load layer 0 weights");
 
     // Load layer 0 (linear attention)
-    let layer = Qwen36LayerWeights::load(&weights, &config, 0, tch::Kind::BFloat16)
-        .expect("load layer 0");
+    let layer =
+        Qwen36LayerWeights::load(&weights, &config, 0, tch::Kind::BFloat16).expect("load layer 0");
 
     // Run forward on CPU (no CUDA in this env)
     let device = tch::Device::Cpu;
-    let input = tch::Tensor::randn(
-        [1, 5, config.hidden_size],
-        (tch::Kind::BFloat16, device),
-    );
+    let input = tch::Tensor::randn([1, 5, config.hidden_size], (tch::Kind::BFloat16, device));
 
-    let output = rustrain_qwen3_6::model::qwen36_layer(&input, &layer, &config, tch::Kind::BFloat16);
+    let output =
+        rustrain_qwen3_6::model::qwen36_layer(&input, &layer, &config, tch::Kind::BFloat16);
 
     assert_eq!(output.size()[0], 1);
     assert_eq!(output.size()[1], 5);
diff --git a/crates/rustrain-qwen3-6/tests/native_cp_attention_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_cp_attention_smoke.cpp
new file mode 100644
index 00000000..9e6908fd
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_cp_attention_smoke.cpp
@@ -0,0 +1,370 @@
+#include <ATen/ATen.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_parallel_nccl_v2(
+    void*, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t);
+extern "C" int32_t qwen36_abort_gradient_accumulation(void*);
+extern "C" int64_t qwen36_export_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" double qwen36_eval_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_micro_step(
+    void*, void*, void*, void*, double, int32_t);
+extern "C" double qwen36_train_multi_lora_selected(
+    void*, void*, void*, void*, const int64_t*, int32_t, int32_t);
+extern "C" void qwen36_free_training_context(void*);
+
+static at::Tensor deterministic(std::initializer_list<int64_t> shape, double scale) {
+    int64_t count = 1;
+    for (int64_t dim : shape) count *= dim;
+    return ((at::arange(count,
+                 at::TensorOptions().device(at::kCUDA).dtype(at::kFloat))
+                 .remainder(19) - 9.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) {
+    return (lhs - rhs).abs().max().item<double>();
+}
+
+int main() {
+    const int rank = std::atoi(std::getenv("RANK"));
+    const int world = std::atoi(std::getenv("WORLD_SIZE"));
+    assert((world == 2 || world == 4) && rank >= 0 && rank < world);
+    const bool ring_test = std::getenv("QWEN36_TEST_CP_RING") &&
+        std::strcmp(std::getenv("QWEN36_TEST_CP_RING"), "0") != 0;
+    qwen36_set_cuda_device(rank);
+    auto set_distributed_env = [&]() {
+        setenv("WORLD_SIZE", std::to_string(world).c_str(), 1);
+        setenv("RANK", std::to_string(rank).c_str(), 1);
+        setenv("TP_SIZE", "1", 1);
+        setenv("CP_SIZE", std::to_string(world).c_str(), 1);
+        setenv("EP_SIZE", "1", 1);
+        setenv("DP_SIZE", "1", 1);
+        setenv("PP_SIZE", "1", 1);
+        setenv("RUSTRAIN_TP_RANK", "0", 1);
+        setenv("RUSTRAIN_CP_RANK", std::to_string(rank).c_str(), 1);
+        setenv("QWEN36_CP_FULL_ATTENTION_KV_GATHER", ring_test ? "0" : "1", 1);
+        setenv("QWEN36_CP_FULL_ATTENTION_RING", ring_test ? "1" : "0", 1);
+    };
+    auto set_reference_env = []() {
+        setenv("WORLD_SIZE", "1", 1);
+        setenv("RANK", "0", 1);
+        setenv("TP_SIZE", "1", 1);
+        setenv("CP_SIZE", "1", 1);
+        setenv("EP_SIZE", "1", 1);
+        setenv("DP_SIZE", "1", 1);
+        setenv("PP_SIZE", "1", 1);
+        setenv("RUSTRAIN_TP_RANK", "0", 1);
+        setenv("RUSTRAIN_CP_RANK", "0", 1);
+        setenv("QWEN36_CP_FULL_ATTENTION_KV_GATHER", "1", 1);
+        setenv("QWEN36_CP_FULL_ATTENTION_RING", "0", 1);
+    };
+
+    constexpr int64_t hidden = 8;
+    constexpr int64_t heads = 4;
+    constexpr int64_t kv_heads = 2;
+    constexpr int64_t head_dim = 2;
+    constexpr int64_t intermediate = 12;
+    constexpr int64_t vocab = 16;
+    constexpr int64_t lora_rank = 4;
+
+    std::vector<at::Tensor> full_weights;
+    full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({2 * heads * head_dim, hidden}, 0.010));
+    full_weights.push_back(at::ones({head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({kv_heads * head_dim, hidden}, 0.012));
+    full_weights.push_back(at::ones({head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({kv_heads * head_dim, hidden}, 0.008));
+    full_weights.push_back(deterministic({hidden, heads * head_dim}, 0.011));
+    full_weights.push_back(deterministic({intermediate, hidden}, 0.009));
+    full_weights.push_back(deterministic({intermediate, hidden}, 0.007));
+    full_weights.push_back(deterministic({hidden, intermediate}, 0.010));
+    for (auto& weight : full_weights) weight.set_requires_grad(false);
+
+    std::vector<at::Tensor> local_weights = full_weights;
+
+    auto embed = deterministic({vocab, hidden}, 0.020);
+    auto final_norm = at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+    auto lm_head = deterministic({vocab, hidden}, 0.015);
+    LayerConfig config{};
+    config.layer_type = 0;
+    config.num_heads = heads;
+    config.num_kv_heads = kv_heads;
+    config.head_dim = head_dim;
+    config.partial_rotary_factor = 1.0;
+    config.rope_theta = 10000.0;
+    config.rms_eps = 1e-5;
+    config.intermediate_size = intermediate;
+
+    const int64_t target_layer = 0;
+    auto local_ptrs = pointers(local_weights);
+    set_distributed_env();
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    void* distributed = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, "q_proj,k_proj,v_proj,o_proj", 0);
+    assert(distributed);
+    assert(qwen36_init_parallel_nccl_v2(
+        distributed, rank, world,
+        0, 1, 0,
+        rank, world, 0,
+        0, 1, 0,
+        0, 1, 0,
+        0, 1, 0) == 0);
+
+    set_reference_env();
+    auto full_ptrs = pointers(full_weights);
+    void* reference = qwen36_create_training_context(
+        full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, "q_proj,k_proj,v_proj,o_proj");
+    assert(reference);
+    set_distributed_env();
+
+    auto q_a = deterministic({lora_rank, hidden}, 0.0020);
+    auto q_b = deterministic({2 * heads * head_dim, lora_rank}, 0.0010);
+    auto k_a = deterministic({lora_rank, hidden}, 0.0018);
+    auto k_b = deterministic({kv_heads * head_dim, lora_rank}, 0.0011);
+    auto v_a = deterministic({lora_rank, hidden}, 0.0016);
+    auto v_b = deterministic({kv_heads * head_dim, lora_rank}, 0.0009);
+    auto o_a = deterministic({lora_rank, heads * head_dim}, 0.0015);
+    auto o_b = deterministic({hidden, lora_rank}, 0.0012);
+    auto local_q_b = q_b.clone();
+    auto local_o_a = o_a.clone();
+    auto local_k_b = k_b.clone();
+    auto local_v_b = v_b.clone();
+    assert(qwen36_set_lora_tensor(distributed, 0, 0, &q_a) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 0, 1, &local_q_b) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 1, 0, &k_a) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 1, 1, &local_k_b) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 2, 0, &v_a) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 2, 1, &local_v_b) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 3, 0, &local_o_a) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 3, 1, &o_b) == 0);
+    assert(qwen36_set_lora_tensor(reference, 0, 0, &q_a) == 0);
+    assert(qwen36_set_lora_tensor(reference, 0, 1, &q_b) == 0);
+    assert(qwen36_set_lora_tensor(reference, 1, 0, &k_a) == 0);
+    assert(qwen36_set_lora_tensor(reference, 1, 1, &k_b) == 0);
+    assert(qwen36_set_lora_tensor(reference, 2, 0, &v_a) == 0);
+    assert(qwen36_set_lora_tensor(reference, 2, 1, &v_b) == 0);
+    assert(qwen36_set_lora_tensor(reference, 3, 0, &o_a) == 0);
+    assert(qwen36_set_lora_tensor(reference, 3, 1, &o_b) == 0);
+
+    auto input_ids = at::tensor({
+        1, 2, 3, 4, 5, 6, 7, 8,
+        9, 10, 11, 12, 13, 14, 0, 0},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({2, 8});
+    auto target_mask = at::tensor({
+        1, 1, 1, 1, 1, 1, 1, 1,
+        1, 1, 1, 1, 1, 1, 0, 0},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)).reshape({2, 8});
+    auto attention_mask = target_mask.to(at::kBool);
+
+    const double distributed_eval = qwen36_eval_step(
+        distributed, &input_ids, &target_mask, &attention_mask);
+    const double reference_eval = qwen36_eval_step(
+        reference, &input_ids, &target_mask, &attention_mask);
+    assert(distributed_eval > 0.0 && reference_eval > 0.0);
+
+    const double distributed_micro = qwen36_train_micro_step(
+        distributed, &input_ids, &target_mask, &attention_mask, 1.0, 0);
+    const double reference_micro = qwen36_train_micro_step(
+        reference, &input_ids, &target_mask, &attention_mask, 1.0, 0);
+    assert(distributed_micro > 0.0 && reference_micro > 0.0);
+    auto* local_q_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(distributed, 0, 1));
+    auto* full_q_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(reference, 0, 1));
+    auto* local_k_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(distributed, 1, 1));
+    auto* full_k_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(reference, 1, 1));
+    auto* local_v_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(distributed, 2, 1));
+    auto* full_v_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(reference, 2, 1));
+    auto* local_o_a_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(distributed, 3, 0));
+    auto* full_o_a_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(reference, 3, 0));
+    assert(local_q_b_grad && full_q_b_grad && local_k_b_grad && full_k_b_grad);
+    assert(local_v_b_grad && full_v_b_grad && local_o_a_grad && full_o_a_grad);
+    const double q_b_grad_diff = max_diff(*local_q_b_grad, *full_q_b_grad);
+    const double o_a_grad_diff = max_diff(*local_o_a_grad, *full_o_a_grad);
+    const double k_b_grad_diff = max_diff(*local_k_b_grad, *full_k_b_grad);
+    const double v_b_grad_diff = max_diff(*local_v_b_grad, *full_v_b_grad);
+    assert(qwen36_abort_gradient_accumulation(distributed) == 0);
+    assert(qwen36_abort_gradient_accumulation(reference) == 0);
+
+    const double distributed_loss = qwen36_train_step(
+        distributed, &input_ids, &target_mask, &attention_mask);
+    const double reference_loss = qwen36_train_step(
+        reference, &input_ids, &target_mask, &attention_mask);
+    assert(distributed_loss > 0.0 && reference_loss > 0.0);
+
+    auto* updated_q_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(distributed, 0));
+    auto* updated_q_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(distributed, 0));
+    auto* updated_k_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(distributed, 1));
+    auto* updated_k_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(distributed, 1));
+    auto* updated_v_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(distributed, 2));
+    auto* updated_v_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(distributed, 2));
+    auto* updated_o_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(distributed, 3));
+    auto* updated_o_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(distributed, 3));
+    auto* reference_q_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(reference, 0));
+    auto* reference_q_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(reference, 0));
+    auto* reference_k_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(reference, 1));
+    auto* reference_k_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(reference, 1));
+    auto* reference_v_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(reference, 2));
+    auto* reference_v_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(reference, 2));
+    auto* reference_o_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(reference, 3));
+    auto* reference_o_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(reference, 3));
+    assert(updated_q_a && updated_q_b && updated_k_a && updated_k_b);
+    assert(updated_v_a && updated_v_b && updated_o_a && updated_o_b);
+    assert(reference_q_a && reference_q_b && reference_k_a && reference_k_b);
+    assert(reference_v_a && reference_v_b && reference_o_a && reference_o_b);
+
+    constexpr int64_t optimizer_count = 14;
+    std::vector<void*> local_m(optimizer_count), local_v(optimizer_count);
+    std::vector<void*> full_m(optimizer_count), full_v(optimizer_count);
+    assert(qwen36_export_optimizer_state(
+        distributed, local_m.data(), local_v.data(), optimizer_count) ==
+        optimizer_count);
+    assert(qwen36_export_optimizer_state(
+        reference, full_m.data(), full_v.data(), optimizer_count) ==
+        optimizer_count);
+    double optimizer_m_diff = 0.0;
+    double optimizer_v_diff = 0.0;
+    double adam_error = 0.0;
+    auto observe_optimizer = [&](int64_t index, const at::Tensor& expected_m,
+                                 const at::Tensor& expected_v,
+                                 const at::Tensor& updated,
+                                 const at::Tensor& before) {
+        auto* m = reinterpret_cast<at::Tensor*>(local_m[index]);
+        auto* v = reinterpret_cast<at::Tensor*>(local_v[index]);
+        assert(m && v);
+        optimizer_m_diff = std::max(optimizer_m_diff, max_diff(*m, expected_m));
+        optimizer_v_diff = std::max(optimizer_v_diff, max_diff(*v, expected_v));
+        auto expected_param = (
+            before.to(at::kFloat) - 1e-3 *
+                (*m / (1.0 - 0.9)) /
+                (((*v / (1.0 - 0.999)).sqrt()) + 1e-8))
+            .to(at::kBFloat16);
+        adam_error = std::max(adam_error, max_diff(updated, expected_param));
+    };
+    auto state = [](std::vector<void*>& tensors, int64_t index) -> at::Tensor& {
+        auto* tensor = reinterpret_cast<at::Tensor*>(tensors[index]);
+        assert(tensor);
+        return *tensor;
+    };
+    observe_optimizer(0, state(full_m, 0), state(full_v, 0), *updated_q_a, q_a);
+    observe_optimizer(1, state(full_m, 1), state(full_v, 1), *updated_q_b, local_q_b);
+    observe_optimizer(2, state(full_m, 2), state(full_v, 2), *updated_k_a, k_a);
+    observe_optimizer(3, state(full_m, 3), state(full_v, 3), *updated_k_b, local_k_b);
+    observe_optimizer(4, state(full_m, 4), state(full_v, 4), *updated_v_a, v_a);
+    observe_optimizer(5, state(full_m, 5), state(full_v, 5), *updated_v_b, local_v_b);
+    observe_optimizer(6, state(full_m, 6), state(full_v, 6), *updated_o_a, local_o_a);
+    observe_optimizer(7, state(full_m, 7), state(full_v, 7), *updated_o_b, o_b);
+    const double q_a_diff = max_diff(*updated_q_a, *reference_q_a);
+    const double q_b_diff = max_diff(*updated_q_b, *reference_q_b);
+    const double o_a_diff = max_diff(*updated_o_a, *reference_o_a);
+    const double o_b_diff = max_diff(*updated_o_b, *reference_o_b);
+    const double k_a_diff = max_diff(*updated_k_a, *reference_k_a);
+    const double k_b_diff = max_diff(*updated_k_b, *reference_k_b);
+    const double v_a_diff = max_diff(*updated_v_a, *reference_v_a);
+    const double v_b_diff = max_diff(*updated_v_b, *reference_v_b);
+    std::printf(
+        "base_cp_attention_smoke rank=%d eval_diff=%0.8e loss_diff=%0.8e q_b_grad_diff=%0.8e k_b_grad_diff=%0.8e v_b_grad_diff=%0.8e o_a_grad_diff=%0.8e m_diff=%0.8e v_diff=%0.8e adam_error=%0.8e q_a_diff=%0.8e q_b_diff=%0.8e k_a_diff=%0.8e k_b_diff=%0.8e v_a_diff=%0.8e v_b_diff=%0.8e o_a_diff=%0.8e o_b_diff=%0.8e\n",
+        rank, std::abs(distributed_eval - reference_eval),
+        std::abs(distributed_loss - reference_loss), q_b_grad_diff,
+        k_b_grad_diff, v_b_grad_diff, o_a_grad_diff, optimizer_m_diff,
+        optimizer_v_diff, adam_error, q_a_diff, q_b_diff, k_a_diff,
+        k_b_diff, v_a_diff, v_b_diff, o_a_diff, o_b_diff);
+    std::fflush(stdout);
+    // The CP path shards token gradients before reducing LoRA gradients.
+    assert(std::abs(distributed_eval - reference_eval) < 5e-3);
+    assert(std::abs(distributed_loss - reference_loss) < 5e-3);
+    // Micro-step accumulators are rank-local until the optimizer boundary.
+    // Their non-zero K/V values exercise the gather backward path; the
+    // synchronized Adam moments below are the CP2-versus-CP1 oracle.
+    assert(std::isfinite(q_b_grad_diff) && std::isfinite(k_b_grad_diff));
+    assert(std::isfinite(v_b_grad_diff) && std::isfinite(o_a_grad_diff));
+    assert(optimizer_m_diff < 5e-5 && optimizer_v_diff < 5e-8);
+    assert(adam_error < 1e-8);
+    const double parameter_threshold = world > 2 ? 3e-3 : 2e-3;
+    assert(std::max({q_a_diff, q_b_diff, k_a_diff, k_b_diff,
+        v_a_diff, v_b_diff, o_a_diff, o_b_diff}) <= parameter_threshold);
+
+    set_distributed_env();
+    if (rank == 1) {
+        setenv("QWEN36_CP_FULL_ATTENTION_KV_GATHER", "0", 1);
+        setenv("QWEN36_CP_FULL_ATTENTION_RING", "0", 1);
+    }
+    const double flag_mismatch = qwen36_eval_step(
+        distributed, &input_ids, &target_mask, &attention_mask);
+    assert(flag_mismatch < 0.0);
+    setenv("QWEN36_CP_FULL_ATTENTION_KV_GATHER", ring_test ? "0" : "1", 1);
+    setenv("QWEN36_CP_FULL_ATTENTION_RING", ring_test ? "1" : "0", 1);
+
+    qwen36_free_training_context(reference);
+    qwen36_free_training_context(distributed);
+
+    std::printf("cp_full_attention_%s_smoke rank=%d flag_mismatch=ok\n",
+        ring_test ? "ring" : "kv_gather", rank);
+    std::fflush(stdout);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_cp_gdn_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_cp_gdn_smoke.cpp
new file mode 100644
index 00000000..f68041aa
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_cp_gdn_smoke.cpp
@@ -0,0 +1,611 @@
+#include <ATen/ATen.h>
+
+#include <algorithm>
+#include <array>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <initializer_list>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" int64_t qwen36_kernel_abi_version();
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_parallel_nccl_v2(
+    void*, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" void* qwen36_get_adapter_optimizer_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t);
+extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t);
+extern "C" int32_t qwen36_train_multi_lora_selected_v3(
+    void*, void*, void*, void*, const int64_t*, int32_t,
+    double*, double*, int32_t);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" int64_t qwen36_export_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" int64_t qwen36_get_step_count(void*);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" double qwen36_eval_step(void*, void*, void*, void*);
+extern "C" double qwen36_eval_step_host_i64(
+    void*, const int64_t*, const int64_t*, const int64_t*, int64_t, int64_t);
+extern "C" void qwen36_set_checkpoint(void*, int32_t, int64_t);
+extern "C" void qwen36_free_training_context(void*);
+
+namespace {
+
+constexpr int64_t kAbiVersion = 31;
+constexpr int64_t kHidden = 32;
+constexpr int64_t kIntermediate = 48;
+constexpr int64_t kKeyHeads = 2;
+constexpr int64_t kValueHeads = 4;
+constexpr int64_t kKeyDim = 128;
+constexpr int64_t kValueDim = 128;
+constexpr int64_t kConvKernel = 4;
+constexpr int64_t kVocab = 64;
+constexpr int64_t kLoraRank = 4;
+constexpr int64_t kPairsPerLayer = 8;
+constexpr int64_t kQSize = kKeyHeads * kKeyDim;
+constexpr int64_t kVSize = kValueHeads * kValueDim;
+constexpr int64_t kQkvSize = 2 * kQSize + kVSize;
+constexpr int64_t kBatch = 2;
+constexpr int64_t kSequence = 8;
+
+constexpr std::array<const char*, 5> kModules = {
+    "in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b", "out_proj"};
+
+static at::Tensor fingerprint(
+    std::initializer_list<int64_t> shape, double scale, int64_t offset
+) {
+    int64_t count = 1;
+    for (const int64_t dim : shape) count *= dim;
+    auto values = at::arange(
+        count, at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    return ((values.add(offset).remainder(97) - 48.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static at::Tensor ones(std::initializer_list<int64_t> shape) {
+    return at::ones(
+        shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+}
+
+static std::vector<at::Tensor> make_weights() {
+    std::vector<at::Tensor> weights;
+    weights.reserve(14);
+    weights.push_back(ones({kHidden}));
+    weights.push_back(ones({kHidden}));
+    auto q = fingerprint({kQSize, kHidden}, 0.00035, 11);
+    auto k = fingerprint({kQSize, kHidden}, 0.00041, 211);
+    auto v = fingerprint({kVSize, kHidden}, 0.00029, 421);
+    weights.push_back(at::cat({q, k, v}, 0).contiguous());
+    weights.push_back(fingerprint({kVSize, kHidden}, 0.00031, 631));
+    weights.push_back(fingerprint({kValueHeads, kHidden}, 0.00043, 719));
+    weights.push_back(fingerprint({kValueHeads, kHidden}, 0.00047, 811));
+    weights.push_back(fingerprint({kValueHeads}, 0.0008, 907));
+    weights.push_back(fingerprint({kValueHeads}, 0.0007, 953));
+    auto q_conv = fingerprint({kQSize, 1, kConvKernel}, 0.0009, 101);
+    auto k_conv = fingerprint({kQSize, 1, kConvKernel}, 0.0011, 307);
+    auto v_conv = fingerprint({kVSize, 1, kConvKernel}, 0.0007, 509);
+    weights.push_back(at::cat({q_conv, k_conv, v_conv}, 0).contiguous());
+    weights.push_back(ones({kValueDim}));
+    weights.push_back(fingerprint({kHidden, kVSize}, 0.00033, 613));
+    weights.push_back(fingerprint({kIntermediate, kHidden}, 0.00045, 701));
+    weights.push_back(fingerprint({kIntermediate, kHidden}, 0.00039, 797));
+    weights.push_back(fingerprint({kHidden, kIntermediate}, 0.00037, 887));
+    return weights;
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static LayerConfig make_config() {
+    LayerConfig config{};
+    config.layer_type = 1;
+    config.num_k_heads = kKeyHeads;
+    config.key_dim = kKeyDim;
+    config.num_v_heads = kValueHeads;
+    config.val_dim = kValueDim;
+    config.conv_kernel = kConvKernel;
+    config.rms_eps = 1e-5;
+    config.intermediate_size = kIntermediate;
+    return config;
+}
+
+struct LoraFixture {
+    int64_t slot;
+    const char* module;
+    at::Tensor a;
+    at::Tensor b;
+};
+
+static int64_t input_size(const char* module) {
+    return std::string(module) == "out_proj" ? kVSize : kHidden;
+}
+
+static int64_t output_size(const char* module) {
+    const std::string name(module);
+    if (name == "in_proj_qkv") return kQkvSize;
+    if (name == "in_proj_z") return kVSize;
+    if (name == "in_proj_a" || name == "in_proj_b") return kValueHeads;
+    assert(name == "out_proj");
+    return kHidden;
+}
+
+static std::vector<LoraFixture> make_lora() {
+    std::vector<LoraFixture> fixtures;
+    fixtures.reserve(kModules.size());
+    for (int64_t slot = 0; slot < static_cast<int64_t>(kModules.size()); ++slot) {
+        const char* module = kModules[slot];
+        fixtures.push_back({
+            slot,
+            module,
+            fingerprint({kLoraRank, input_size(module)}, 0.0007, 1201 + slot * 31),
+            fingerprint({output_size(module), kLoraRank}, 0.0006, 1601 + slot * 37),
+        });
+    }
+    return fixtures;
+}
+
+static void set_dynamic_lora(
+    void* context, int64_t adapter_id,
+    std::vector<LoraFixture>& fixtures
+) {
+    for (auto& fixture : fixtures) {
+        assert(qwen36_set_adapter_lora_tensor(
+            context, adapter_id, 0, fixture.module, 0, &fixture.a) == 0);
+        assert(qwen36_set_adapter_lora_tensor(
+            context, adapter_id, 0, fixture.module, 1, &fixture.b) == 0);
+    }
+}
+
+struct Batch {
+    at::Tensor ids;
+    at::Tensor targets;
+    at::Tensor attention;
+};
+
+static Batch make_right_padded_batch() {
+    std::vector<int64_t> ids(kBatch * kSequence);
+    std::vector<float> targets(kBatch * kSequence, 1.0f);
+    std::vector<uint8_t> attention(kBatch * kSequence, 1);
+    for (int64_t batch = 0; batch < kBatch; ++batch) {
+        for (int64_t token = 0; token < kSequence; ++token) {
+            ids[batch * kSequence + token] =
+                1 + (batch * 17 + token * 5) % (kVocab - 1);
+        }
+    }
+    for (int64_t token = 6; token < kSequence; ++token) {
+        targets[kSequence + token] = 0.0f;
+        attention[kSequence + token] = 0;
+    }
+    auto ids_tensor = at::from_blob(
+        ids.data(), {kBatch, kSequence},
+        at::TensorOptions().device(at::kCPU).dtype(at::kLong)).clone().to(at::kCUDA);
+    auto targets_tensor = at::from_blob(
+        targets.data(), {kBatch, kSequence},
+        at::TensorOptions().device(at::kCPU).dtype(at::kFloat)).clone().to(at::kCUDA);
+    auto attention_tensor = at::from_blob(
+        attention.data(), {kBatch, kSequence},
+        at::TensorOptions().device(at::kCPU).dtype(at::kByte)).clone()
+        .to(at::kCUDA).to(at::kBool);
+    return {std::move(ids_tensor), std::move(targets_tensor),
+        std::move(attention_tensor)};
+}
+
+static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) {
+    assert(lhs.sizes() == rhs.sizes());
+    return (lhs.to(at::kFloat) - rhs.to(at::kFloat))
+        .abs().max().item<double>();
+}
+
+static void set_parallel_environment(int rank, int world, int cp_rank, int cp_size) {
+    setenv("WORLD_SIZE", std::to_string(world).c_str(), 1);
+    setenv("RANK", std::to_string(rank).c_str(), 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("CP_SIZE", std::to_string(cp_size).c_str(), 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("PP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    setenv("RUSTRAIN_CP_RANK", std::to_string(cp_rank).c_str(), 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    setenv("RUSTRAIN_PP_RANK", "0", 1);
+}
+
+static void* create_context(
+    std::vector<at::Tensor>& weights, at::Tensor& embedding,
+    at::Tensor& final_norm, at::Tensor& lm_head, LayerConfig& config,
+    bool extended
+) {
+    auto weight_ptrs = pointers(weights);
+    const int64_t target_layer = 0;
+    if (extended) {
+        return qwen36_create_training_context_ex(
+            weight_ptrs.data(), weight_ptrs.size(), &embedding, &final_norm,
+            &lm_head, &config, 1, static_cast<int32_t>(at::kBFloat16),
+            1.0, 1e-3, 0.9, 0.999, 1e-8, kVocab, 1e-5, kLoraRank,
+            &target_layer, 1,
+            "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj", 0);
+    }
+    return qwen36_create_training_context(
+        weight_ptrs.data(), weight_ptrs.size(), &embedding, &final_norm,
+        &lm_head, &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, kVocab, 1e-5, kLoraRank,
+        &target_layer, 1,
+        "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj");
+}
+
+}  // namespace
+
+int main() {
+    assert(qwen36_kernel_abi_version() == kAbiVersion);
+    const int rank = std::atoi(std::getenv("RANK"));
+    const int world = std::atoi(std::getenv("WORLD_SIZE"));
+    const int local_rank = std::atoi(std::getenv("LOCAL_RANK"));
+    assert(world == 2 && (rank == 0 || rank == 1));
+    qwen36_set_cuda_device(local_rank);
+    unsetenv("QWEN36_SEQ_CHUNK");
+    auto weights = make_weights();
+    auto embedding = fingerprint({kVocab, kHidden}, 0.0021, 2001);
+    auto final_norm = ones({kHidden});
+    auto lm_head = fingerprint({kVocab, kHidden}, 0.0017, 2201);
+    auto config = make_config();
+
+    set_parallel_environment(rank, world, rank, 2);
+    void* distributed = create_context(
+        weights, embedding, final_norm, lm_head, config, true);
+    assert(distributed);
+    auto* initial_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_a(distributed, 0));
+    auto* initial_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_b(distributed, 0));
+    assert(initial_a && initial_b);
+    initial_a->detach().fill_(rank == 0 ? 0.125 : -0.25);
+    initial_b->detach().fill_(rank == 0 ? 0.25 : -0.5);
+
+    set_parallel_environment(0, 1, 0, 1);
+    void* reference = create_context(
+        weights, embedding, final_norm, lm_head, config, false);
+    assert(reference);
+
+    set_parallel_environment(rank, world, rank, 2);
+    assert(qwen36_init_parallel_nccl_v2(
+        distributed, rank, world,
+        0, 1, 0,
+        rank, 2, 0,
+        0, 1, 0,
+        0, 1, 0,
+        0, 1, 0) == 0);
+    assert((initial_a->to(at::kFloat) - 0.125).abs().max().item<float>() == 0.0f);
+    assert((initial_b->to(at::kFloat) - 0.25).abs().max().item<float>() == 0.0f);
+    qwen36_set_checkpoint(distributed, 1, 1);
+    qwen36_set_checkpoint(reference, 1, 1);
+
+    auto fixtures = make_lora();
+    for (auto& fixture : fixtures) {
+        assert(qwen36_set_lora_tensor(
+            distributed, fixture.slot, 0, &fixture.a) == 0);
+        assert(qwen36_set_lora_tensor(
+            distributed, fixture.slot, 1, &fixture.b) == 0);
+        assert(qwen36_set_lora_tensor(
+            reference, fixture.slot, 0, &fixture.a) == 0);
+        assert(qwen36_set_lora_tensor(
+            reference, fixture.slot, 1, &fixture.b) == 0);
+    }
+
+    auto batch = make_right_padded_batch();
+    const double distributed_eval_loss = qwen36_eval_step(
+        distributed, &batch.ids, &batch.targets, &batch.attention);
+    const double reference_eval_loss = qwen36_eval_step(
+        reference, &batch.ids, &batch.targets, &batch.attention);
+    assert(std::isfinite(distributed_eval_loss) &&
+        std::isfinite(reference_eval_loss));
+    assert(std::abs(distributed_eval_loss - reference_eval_loss) < 1e-2);
+
+    const double distributed_loss = qwen36_train_step(
+        distributed, &batch.ids, &batch.targets, &batch.attention);
+    const double reference_loss = qwen36_train_step(
+        reference, &batch.ids, &batch.targets, &batch.attention);
+    assert(std::isfinite(distributed_loss) && std::isfinite(reference_loss));
+
+    constexpr int64_t optimizer_count = 2 * kPairsPerLayer;
+    std::vector<void*> distributed_m(optimizer_count), distributed_v(optimizer_count);
+    std::vector<void*> reference_m(optimizer_count), reference_v(optimizer_count);
+    assert(qwen36_export_optimizer_state(
+        distributed, distributed_m.data(), distributed_v.data(),
+        optimizer_count) == optimizer_count);
+    assert(qwen36_export_optimizer_state(
+        reference, reference_m.data(), reference_v.data(),
+        optimizer_count) == optimizer_count);
+
+    double max_a_diff = 0.0;
+    double max_b_diff = 0.0;
+    double max_m_diff = 0.0;
+    double max_v_diff = 0.0;
+    double max_update = 0.0;
+    for (const auto& fixture : fixtures) {
+        auto* distributed_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(distributed, fixture.slot));
+        auto* distributed_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(distributed, fixture.slot));
+        auto* reference_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(reference, fixture.slot));
+        auto* reference_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(reference, fixture.slot));
+        assert(distributed_a && distributed_b && reference_a && reference_b);
+        max_a_diff = std::max(max_a_diff, max_diff(*distributed_a, *reference_a));
+        max_b_diff = std::max(max_b_diff, max_diff(*distributed_b, *reference_b));
+        max_update = std::max({max_update,
+            max_diff(*distributed_a, fixture.a),
+            max_diff(*distributed_b, fixture.b)});
+
+        for (int is_b = 0; is_b < 2; ++is_b) {
+            const int64_t state_index = 2 * fixture.slot + is_b;
+            auto* dist_m = reinterpret_cast<at::Tensor*>(distributed_m[state_index]);
+            auto* dist_v = reinterpret_cast<at::Tensor*>(distributed_v[state_index]);
+            auto* ref_m = reinterpret_cast<at::Tensor*>(reference_m[state_index]);
+            auto* ref_v = reinterpret_cast<at::Tensor*>(reference_v[state_index]);
+            assert(dist_m && dist_v && ref_m && ref_v);
+            max_m_diff = std::max(max_m_diff, max_diff(*dist_m, *ref_m));
+            max_v_diff = std::max(max_v_diff, max_diff(*dist_v, *ref_v));
+        }
+    }
+
+    const double loss_diff = std::abs(distributed_loss - reference_loss);
+    const int64_t distributed_step = qwen36_get_step_count(distributed);
+    const int64_t reference_step = qwen36_get_step_count(reference);
+    std::printf(
+        "native_cp_gdn_smoke rank=%d eval_loss_diff=%0.8e "
+        "loss_diff=%0.8e a_diff=%0.8e "
+        "b_diff=%0.8e m_diff=%0.8e v_diff=%0.8e max_update=%0.8e "
+        "step=%ld reference_step=%ld\n",
+        rank, std::abs(distributed_eval_loss - reference_eval_loss),
+        loss_diff, max_a_diff, max_b_diff, max_m_diff, max_v_diff,
+        max_update, static_cast<long>(distributed_step),
+        static_cast<long>(reference_step));
+    std::fflush(stdout);
+
+    assert(loss_diff < 1e-2);
+    assert(max_a_diff <= 2.1e-3 && max_b_diff <= 2.1e-3);
+    assert(max_m_diff < 5e-4 && max_v_diff < 5e-7);
+    assert(max_update > 0.0);
+    assert(distributed_step == 1 && reference_step == 1);
+
+    const char* fused_cp_exchange_env =
+        getenv("QWEN36_GDN_FUSED_CP_EXCHANGE");
+    const bool had_fused_cp_exchange_env = fused_cp_exchange_env != nullptr;
+    const std::string saved_fused_cp_exchange =
+        fused_cp_exchange_env ? fused_cp_exchange_env : "";
+    const bool fused_cp_exchange_enabled = fused_cp_exchange_env &&
+        std::string(fused_cp_exchange_env) != "0";
+    if (rank == 1) {
+        setenv("QWEN36_GDN_FUSED_CP_EXCHANGE",
+            fused_cp_exchange_enabled ? "0" : "1", 1);
+    }
+    assert(qwen36_eval_step(
+        distributed, &batch.ids, &batch.targets, &batch.attention) < 0.0);
+    if (rank == 1) {
+        if (had_fused_cp_exchange_env) {
+            setenv("QWEN36_GDN_FUSED_CP_EXCHANGE",
+                saved_fused_cp_exchange.c_str(), 1);
+        } else {
+            unsetenv("QWEN36_GDN_FUSED_CP_EXCHANGE");
+        }
+    }
+
+    if (rank == 1) setenv("QWEN36_SEQ_CHUNK", "4", 1);
+    assert(qwen36_eval_step(
+        distributed, &batch.ids, &batch.targets, &batch.attention) < 0.0);
+    unsetenv("QWEN36_SEQ_CHUNK");
+    if (rank == 1) batch.ids[0][0].fill_(batch.ids[0][0].item<int64_t>() + 1);
+    assert(qwen36_eval_step(
+        distributed, &batch.ids, &batch.targets, &batch.attention) < 0.0);
+
+    std::vector<int64_t> host_ids(kBatch * kSequence, 1);
+    std::vector<int64_t> host_targets(kBatch * kSequence, 1);
+    std::vector<int64_t> host_attention(kBatch * kSequence, 1);
+    const int64_t* host_input = rank == 1 ? nullptr : host_ids.data();
+    assert(qwen36_eval_step_host_i64(
+        distributed, host_input, host_targets.data(), host_attention.data(),
+        kBatch, kSequence) < 0.0);
+
+    qwen36_free_training_context(reference);
+    qwen36_free_training_context(distributed);
+
+    // Dynamic CP keeps the full tenant batch descriptor replicated while
+    // each rank owns one sequence/head contribution. Two selected tenants
+    // use different token counts; a third tenant must remain bitwise idle.
+    auto dynamic_batch = make_right_padded_batch();
+    set_parallel_environment(rank, world, rank, 2);
+    void* dynamic_distributed = create_context(
+        weights, embedding, final_norm, lm_head, config, true);
+    assert(dynamic_distributed);
+    assert(qwen36_init_parallel_nccl_v2(
+        dynamic_distributed, rank, world,
+        0, 1, 0,
+        rank, 2, 0,
+        0, 1, 0,
+        0, 1, 0,
+        0, 1, 0) == 0);
+
+    set_parallel_environment(0, 1, 0, 1);
+    void* dynamic_reference = create_context(
+        weights, embedding, final_norm, lm_head, config, false);
+    assert(dynamic_reference);
+    set_parallel_environment(rank, world, rank, 2);
+
+    const int64_t dynamic_target_layer = 0;
+    constexpr const char* dynamic_targets =
+        "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj";
+    std::array<int64_t, 3> distributed_adapters{};
+    std::array<int64_t, 3> reference_adapters{};
+    for (int64_t index = 0; index < 3; ++index) {
+        distributed_adapters[index] = qwen36_add_lora(
+            dynamic_distributed, kLoraRank, kLoraRank,
+            &dynamic_target_layer, 1, dynamic_targets);
+        reference_adapters[index] = qwen36_add_lora(
+            dynamic_reference, kLoraRank, kLoraRank,
+            &dynamic_target_layer, 1, dynamic_targets);
+        assert(distributed_adapters[index] > 0 &&
+            reference_adapters[index] > 0);
+        set_dynamic_lora(
+            dynamic_distributed, distributed_adapters[index], fixtures);
+        set_dynamic_lora(
+            dynamic_reference, reference_adapters[index], fixtures);
+    }
+
+    double distributed_dynamic_loss = -1.0;
+    double reference_dynamic_loss = -1.0;
+    std::array<double, 2> distributed_tenant_losses{-1.0, -1.0};
+    std::array<double, 2> reference_tenant_losses{-1.0, -1.0};
+    assert(qwen36_train_multi_lora_selected_v3(
+        dynamic_distributed, &dynamic_batch.ids, &dynamic_batch.targets,
+        &dynamic_batch.attention, distributed_adapters.data(), 2,
+        &distributed_dynamic_loss, distributed_tenant_losses.data(), 2) == 0);
+    assert(qwen36_train_multi_lora_selected_v3(
+        dynamic_reference, &dynamic_batch.ids, &dynamic_batch.targets,
+        &dynamic_batch.attention, reference_adapters.data(), 2,
+        &reference_dynamic_loss, reference_tenant_losses.data(), 2) == 0);
+
+    auto adapter_tensor = [](void* context, int64_t adapter_id,
+                             const LoraFixture& fixture, bool is_b) {
+        auto* tensor = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_lora_tensor(
+                context, adapter_id, 0, fixture.module, is_b ? 1 : 0));
+        assert(tensor);
+        return tensor;
+    };
+    auto optimizer_tensor = [](void* context, int64_t adapter_id,
+                               const LoraFixture& fixture,
+                               bool is_b, bool is_v) {
+        auto* tensor = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_optimizer_tensor(
+                context, adapter_id, 0, fixture.module,
+                is_b ? 1 : 0, is_v ? 1 : 0));
+        assert(tensor && tensor->scalar_type() == at::kFloat);
+        return tensor;
+    };
+
+    double dynamic_param_diff = 0.0;
+    double dynamic_m_diff = 0.0;
+    double dynamic_v_diff = 0.0;
+    double dynamic_update = 0.0;
+    double unselected_diff = 0.0;
+    for (int64_t adapter_index = 0; adapter_index < 2; ++adapter_index) {
+        for (const auto& fixture : fixtures) {
+            for (int is_b = 0; is_b < 2; ++is_b) {
+                auto* distributed_param = adapter_tensor(
+                    dynamic_distributed,
+                    distributed_adapters[adapter_index], fixture, is_b != 0);
+                auto* reference_param = adapter_tensor(
+                    dynamic_reference,
+                    reference_adapters[adapter_index], fixture, is_b != 0);
+                dynamic_param_diff = std::max(dynamic_param_diff,
+                    max_diff(*distributed_param, *reference_param));
+                dynamic_update = std::max(dynamic_update,
+                    max_diff(*distributed_param,
+                        is_b ? fixture.b : fixture.a));
+                auto* distributed_m = optimizer_tensor(
+                    dynamic_distributed,
+                    distributed_adapters[adapter_index], fixture,
+                    is_b != 0, false);
+                auto* reference_m = optimizer_tensor(
+                    dynamic_reference,
+                    reference_adapters[adapter_index], fixture,
+                    is_b != 0, false);
+                auto* distributed_v = optimizer_tensor(
+                    dynamic_distributed,
+                    distributed_adapters[adapter_index], fixture,
+                    is_b != 0, true);
+                auto* reference_v = optimizer_tensor(
+                    dynamic_reference,
+                    reference_adapters[adapter_index], fixture,
+                    is_b != 0, true);
+                dynamic_m_diff = std::max(dynamic_m_diff,
+                    max_diff(*distributed_m, *reference_m));
+                dynamic_v_diff = std::max(dynamic_v_diff,
+                    max_diff(*distributed_v, *reference_v));
+            }
+        }
+    }
+    for (const auto& fixture : fixtures) {
+        for (int is_b = 0; is_b < 2; ++is_b) {
+            auto* parameter = adapter_tensor(
+                dynamic_distributed, distributed_adapters[2], fixture,
+                is_b != 0);
+            unselected_diff = std::max(unselected_diff,
+                max_diff(*parameter, is_b ? fixture.b : fixture.a));
+            for (int is_v = 0; is_v < 2; ++is_v) {
+                auto* state = optimizer_tensor(
+                    dynamic_distributed, distributed_adapters[2], fixture,
+                    is_b != 0, is_v != 0);
+                unselected_diff = std::max(unselected_diff,
+                    state->abs().max().item<double>());
+            }
+        }
+    }
+
+    const double dynamic_loss_diff =
+        std::abs(distributed_dynamic_loss - reference_dynamic_loss);
+    const double tenant_loss_diff = std::max(
+        std::abs(distributed_tenant_losses[0] - reference_tenant_losses[0]),
+        std::abs(distributed_tenant_losses[1] - reference_tenant_losses[1]));
+    std::printf(
+        "native_cp_gdn_dynamic rank=%d loss_diff=%0.8e "
+        "tenant_loss_diff=%0.8e param_diff=%0.8e m_diff=%0.8e "
+        "v_diff=%0.8e update=%0.8e unselected_diff=%0.8e\n",
+        rank, dynamic_loss_diff, tenant_loss_diff, dynamic_param_diff,
+        dynamic_m_diff, dynamic_v_diff, dynamic_update, unselected_diff);
+    std::fflush(stdout);
+    assert(dynamic_loss_diff < 1e-2 && tenant_loss_diff < 1e-2);
+    assert(dynamic_param_diff <= 2.1e-3);
+    assert(dynamic_m_diff < 5e-4 && dynamic_v_diff < 5e-7);
+    assert(dynamic_update > 0.0 && unselected_diff == 0.0);
+    assert(qwen36_get_adapter_step_count(
+        dynamic_distributed, distributed_adapters[0]) == 1);
+    assert(qwen36_get_adapter_step_count(
+        dynamic_distributed, distributed_adapters[1]) == 1);
+    assert(qwen36_get_adapter_step_count(
+        dynamic_distributed, distributed_adapters[2]) == 0);
+
+    qwen36_free_training_context(dynamic_reference);
+    qwen36_free_training_context(dynamic_distributed);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_ep_bench.cpp b/crates/rustrain-qwen3-6/tests/native_ep_bench.cpp
new file mode 100644
index 00000000..a0a67a81
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_ep_bench.cpp
@@ -0,0 +1,228 @@
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <initializer_list>
+#include <numeric>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" int32_t qwen36_init_nccl(void*);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" void qwen36_free_training_context(void*);
+extern "C" double qwen36_parallel_max_double(void*, double);
+
+static int env_int(const char* name, int fallback) {
+    const char* value = std::getenv(name);
+    return value ? std::max(1, std::atoi(value)) : fallback;
+}
+
+static bool env_enabled(const char* name, bool fallback = false) {
+    const char* value = std::getenv(name);
+    if (!value || value[0] == '\0') return fallback;
+    return std::strcmp(value, "0") != 0 && std::strcmp(value, "false") != 0;
+}
+
+static at::Tensor cuda_rand(std::initializer_list<int64_t> shape) {
+    return at::randn(shape,
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+}
+
+static std::vector<void*> tensor_ptrs(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static double percentile(std::vector<double> values, double q) {
+    std::sort(values.begin(), values.end());
+    const double pos = q * static_cast<double>(values.size() - 1);
+    const size_t lo = static_cast<size_t>(pos);
+    const size_t hi = std::min(lo + 1, values.size() - 1);
+    const double frac = pos - static_cast<double>(lo);
+    return values[lo] * (1.0 - frac) + values[hi] * frac;
+}
+
+int main() {
+    const int rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0");
+    const int world = std::atoi(
+        std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1");
+    const int local_rank = std::atoi(
+        std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0");
+    const bool a2a = env_enabled("QWEN36_EP_A2A");
+    const bool sharded = a2a && env_enabled("QWEN36_EP_A2A_SHARDED");
+    const bool packed = env_enabled("QWEN36_EP_A2A_PACKED", true);
+    const int seq = env_int("BENCH_SEQ", 128);
+    const int hidden = env_int("BENCH_HIDDEN", 256);
+    const int experts = env_int("BENCH_EXPERTS", 8);
+    const int intermediate = env_int("BENCH_INTERMEDIATE", 256);
+    const int warmup = env_int("BENCH_WARMUP", 2);
+    const int iters = env_int("BENCH_ITERS", 10);
+    const bool gpu_metadata = env_enabled(
+        "QWEN36_EP_A2A_GPU_METADATA",
+        world >= 4 && seq * (packed ? 2 : 1) >= 512);
+    assert(world >= 2 && rank >= 0 && rank < world);
+    assert(hidden % 8 == 0 && experts % world == 0);
+    assert(!sharded || a2a);
+    c10::cuda::CUDAGuard guard(local_rank);
+    at::manual_seed(1234);
+
+    const int local_experts = experts / world;
+    const int64_t expert_start = rank * local_experts;
+    const int vocab = std::max(1024, hidden * 4);
+    const int head_dim = hidden;
+
+    // Build one deterministic global model, then narrow only expert tensors.
+    std::vector<at::Tensor> global_weights;
+    global_weights.push_back(cuda_rand({hidden}));
+    global_weights.push_back(cuda_rand({hidden}));
+    global_weights.push_back(cuda_rand({2 * head_dim, hidden}));
+    global_weights.push_back(cuda_rand({head_dim}));
+    global_weights.push_back(cuda_rand({head_dim, hidden}));
+    global_weights.push_back(cuda_rand({head_dim}));
+    global_weights.push_back(cuda_rand({head_dim, hidden}));
+    global_weights.push_back(cuda_rand({hidden, head_dim}));
+    global_weights.push_back(cuda_rand({experts, hidden}));
+    global_weights.push_back(cuda_rand({1, hidden}));
+    global_weights.push_back(cuda_rand({intermediate, hidden}));
+    global_weights.push_back(cuda_rand({intermediate, hidden}));
+    global_weights.push_back(cuda_rand({hidden, intermediate}));
+    global_weights.push_back(cuda_rand({experts, 2 * intermediate, hidden}));
+    global_weights.push_back(cuda_rand({experts, hidden, intermediate}));
+    for (auto& weight : global_weights) weight.set_requires_grad(false);
+
+    std::vector<at::Tensor> local_weights = global_weights;
+    local_weights[13] = global_weights[13]
+        .narrow(0, expert_start, local_experts).contiguous();
+    local_weights[14] = global_weights[14]
+        .narrow(0, expert_start, local_experts).contiguous();
+    auto weights = tensor_ptrs(local_weights);
+    auto embed = cuda_rand({vocab, hidden});
+    auto final_norm = cuda_rand({hidden});
+    auto lm_head = cuda_rand({vocab, hidden});
+    embed.set_requires_grad(false);
+    final_norm.set_requires_grad(false);
+    lm_head.set_requires_grad(false);
+
+    LayerConfig config{};
+    config.layer_type = 0;
+    config.num_heads = 1;
+    config.num_kv_heads = 1;
+    config.head_dim = head_dim;
+    config.partial_rotary_factor = 1.0;
+    config.rope_theta = 10000.0;
+    config.rms_eps = 1e-5;
+    config.num_experts = experts;
+    config.top_k = 2;
+    config.moe_intermediate = intermediate;
+    config.expert_start = expert_start;
+    config.expert_count = local_experts;
+    config.norm_topk_prob = 1;
+
+    const int64_t target_layer = 0;
+    const char* targets = "experts_gate_up_proj,experts_down_proj";
+    void* ctx = qwen36_create_training_context(
+        weights.data(), static_cast<int64_t>(weights.size()),
+        &embed, &final_norm, &lm_head, &config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, 8,
+        &target_layer, 1, targets);
+    assert(ctx);
+    assert(qwen36_init_nccl(ctx) == 0);
+
+    std::vector<int64_t> ids(seq);
+    for (int i = 0; i < seq; ++i) {
+        const int offset = sharded ? rank * seq : 0;
+        ids[i] = (offset + i + 1) % vocab;
+    }
+    auto input_ids = at::from_blob(ids.data(), {1, seq},
+        at::TensorOptions().device(at::kCPU).dtype(at::kLong)).clone().to(at::kCUDA);
+    auto target_mask = at::ones({1, seq},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    auto attention_mask = at::ones({1, seq},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+
+    for (int i = 0; i < warmup; ++i) {
+        (void)qwen36_train_step(ctx, &input_ids, &target_mask, &attention_mask);
+    }
+    cudaDeviceSynchronize();
+
+    cudaEvent_t start = nullptr, stop = nullptr;
+    cudaEventCreate(&start);
+    cudaEventCreate(&stop);
+    std::vector<double> times;
+    times.reserve(iters);
+    double last_loss = 0.0;
+    for (int i = 0; i < iters; ++i) {
+        cudaEventRecord(start, c10::cuda::getCurrentCUDAStream().stream());
+        last_loss = qwen36_train_step(
+            ctx, &input_ids, &target_mask, &attention_mask);
+        cudaEventRecord(stop, c10::cuda::getCurrentCUDAStream().stream());
+        cudaEventSynchronize(stop);
+        float elapsed_ms = 0.0f;
+        cudaEventElapsedTime(&elapsed_ms, start, stop);
+        const double world_max_ms = qwen36_parallel_max_double(
+            ctx, static_cast<double>(elapsed_ms));
+        assert(world_max_ms > 0.0 && std::isfinite(world_max_ms));
+        times.push_back(world_max_ms);
+    }
+    cudaEventDestroy(start);
+    cudaEventDestroy(stop);
+
+    const double mean = std::accumulate(times.begin(), times.end(), 0.0) /
+        static_cast<double>(times.size());
+    double variance = 0.0;
+    for (const double value : times) variance += (value - mean) * (value - mean);
+    variance /= static_cast<double>(times.size());
+    size_t free_bytes = 0, total_bytes = 0;
+    cudaMemGetInfo(&free_bytes, &total_bytes);
+    const double local_tokens = static_cast<double>(std::max(seq - 1, 1));
+    const double processed_tokens = local_tokens * world;
+    // Legacy EP replicates the input batch on every rank; only sharded A2A
+    // represents distinct global samples in this synthetic harness.
+    const double unique_tokens = sharded ? processed_tokens : local_tokens;
+    const double median_seconds = percentile(times, 0.5) / 1000.0;
+    const double processed_tokens_per_sec = processed_tokens / median_seconds;
+    const double unique_tokens_per_sec = unique_tokens / median_seconds;
+    std::printf(
+        "native_qwen36_ep_bench rank=%d world=%d a2a=%d sharded=%d packed=%d "
+        "gpu_metadata=%d "
+        "seq=%d hidden=%d experts=%d intermediate=%d warmup=%d iters=%d "
+        "local_tokens=%.0f processed_tokens=%.0f unique_tokens=%.0f last_loss=%.8f "
+        "timing_scope=world_max step_ms_mean=%.4f step_ms_p50=%.4f "
+        "step_ms_p95=%.4f step_ms_std=%.4f "
+        "processed_tokens_per_sec=%.2f unique_tokens_per_sec=%.2f "
+        "free_mem_gib=%.3f\n",
+        rank, world, a2a, sharded, packed, gpu_metadata,
+        seq, hidden, experts, intermediate,
+        warmup, iters, local_tokens, processed_tokens, unique_tokens, last_loss,
+        mean, percentile(times, 0.5), percentile(times, 0.95),
+        std::sqrt(variance),
+        processed_tokens_per_sec, unique_tokens_per_sec,
+        static_cast<double>(free_bytes) / (1024.0 * 1024.0 * 1024.0));
+    std::fflush(stdout);
+    qwen36_free_training_context(ctx);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp
new file mode 100644
index 00000000..3f426c0a
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp
@@ -0,0 +1,659 @@
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" int32_t qwen36_init_nccl(void*);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" void* qwen36_get_adapter_optimizer_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t);
+extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t);
+extern "C" double qwen36_train_multi_lora(
+    void*, void*, void*, void*, int32_t, int32_t);
+extern "C" int64_t qwen36_get_lora_count(void*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" int64_t qwen36_export_optimizer_state(void*, void**, void**, int64_t);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" void qwen36_free_training_context(void*);
+
+static at::Tensor cuda_rand(std::initializer_list<int64_t> shape) {
+    return at::randn(shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+}
+
+static std::vector<void*> tensor_ptrs(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> ptrs;
+    ptrs.reserve(tensors.size());
+    for (auto& tensor : tensors) ptrs.push_back(&tensor);
+    return ptrs;
+}
+
+static double max_abs_diff(const at::Tensor& actual, const at::Tensor& expected) {
+    assert(actual.sizes() == expected.sizes());
+    return (actual.to(at::kFloat) - expected.to(at::kFloat))
+        .abs().max().item<double>();
+}
+
+static double update_norm(const at::Tensor& after, const at::Tensor& before) {
+    return (after.to(at::kFloat) - before.to(at::kFloat))
+        .abs().sum().item<double>();
+}
+
+static double first_adam_step_diff(
+    const at::Tensor& actual,
+    const at::Tensor& before,
+    const at::Tensor& first_m,
+    const at::Tensor& first_v
+) {
+    constexpr double lr = 1e-3;
+    constexpr double beta1 = 0.9;
+    constexpr double beta2 = 0.999;
+    constexpr double eps = 1e-8;
+    auto m_hat = first_m.to(at::kFloat) / (1.0 - beta1);
+    auto v_hat = first_v.to(at::kFloat) / (1.0 - beta2);
+    auto expected = (before.to(at::kFloat) -
+        lr * m_hat / (v_hat.sqrt() + eps)).to(actual.scalar_type());
+    return max_abs_diff(actual, expected);
+}
+
+struct ExpertLora {
+    at::Tensor* gate_up_a;
+    at::Tensor* gate_up_b;
+    at::Tensor* down_a;
+    at::Tensor* down_b;
+};
+
+static ExpertLora get_expert_lora(void* ctx) {
+    assert(qwen36_get_lora_count(ctx) == 9);
+    ExpertLora lora{
+        reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(ctx, 7)),
+        reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(ctx, 7)),
+        reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(ctx, 8)),
+        reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(ctx, 8)),
+    };
+    assert(lora.gate_up_a && lora.gate_up_b && lora.down_a && lora.down_b);
+    return lora;
+}
+
+static void set_expert_lora(
+    void* ctx,
+    at::Tensor& gate_up_a,
+    at::Tensor& gate_up_b,
+    at::Tensor& down_a,
+    at::Tensor& down_b
+) {
+    assert(qwen36_set_lora_tensor(ctx, 7, 0, &gate_up_a) == 0);
+    assert(qwen36_set_lora_tensor(ctx, 7, 1, &gate_up_b) == 0);
+    assert(qwen36_set_lora_tensor(ctx, 8, 0, &down_a) == 0);
+    assert(qwen36_set_lora_tensor(ctx, 8, 1, &down_b) == 0);
+}
+
+int main() {
+    const int rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0");
+    const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1");
+    const int local_rank = std::atoi(
+        std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0");
+    const int a2a = std::getenv("QWEN36_EP_A2A") &&
+        std::strcmp(std::getenv("QWEN36_EP_A2A"), "0") != 0;
+    const int sharded = a2a && std::getenv("QWEN36_EP_A2A_SHARDED") &&
+        std::strcmp(std::getenv("QWEN36_EP_A2A_SHARDED"), "0") != 0;
+    const bool dynamic_only = std::getenv("QWEN36_DYNAMIC_ONLY") &&
+        std::strcmp(std::getenv("QWEN36_DYNAMIC_ONLY"), "0") != 0;
+    // Replicated A2A receives duplicate source rows in rank order, so its BF16
+    // accumulation is not bit-identical to the full-expert reference. Sharded
+    // A2A uses distinct rows; its optimizer state matches closely, while a
+    // near-boundary update can land in the adjacent BF16 parameter bin. Keep
+    // the legacy threshold strict, bound each BF16 case separately, and retain
+    // the exact standard-Adam oracle below for all modes.
+    const double param_tol = sharded ? 2e-3 : (a2a ? 2e-4 : 1e-5);
+    const double m_tol = sharded ? 5e-4 : (a2a ? 5e-3 : 1e-5);
+    const double v_tol = a2a ? 1e-5 : 1e-6;
+    assert(world == 2 && rank >= 0 && rank < world);
+    assert(!std::getenv("TP_SIZE") || std::atoi(std::getenv("TP_SIZE")) == 1);
+    c10::cuda::CUDAGuard guard(local_rank);
+
+    // Every process deterministically creates the same global model. The EP
+    // context receives a distinct contiguous expert slice from that model.
+    at::manual_seed(100);
+    constexpr int64_t hidden = 16;
+    constexpr int64_t vocab = 8;
+    constexpr int64_t experts = 2;
+    constexpr int64_t head_dim = 8;
+    constexpr int64_t intermediate = 8;
+    constexpr int64_t lora_rank = 4;
+
+    std::vector<at::Tensor> global_weights;
+    global_weights.push_back(cuda_rand({hidden}));
+    global_weights.push_back(cuda_rand({hidden}));
+    global_weights.push_back(cuda_rand({2 * head_dim, hidden}));
+    global_weights.push_back(cuda_rand({head_dim}));
+    global_weights.push_back(cuda_rand({head_dim, hidden}));
+    global_weights.push_back(cuda_rand({head_dim}));
+    global_weights.push_back(cuda_rand({head_dim, hidden}));
+    global_weights.push_back(cuda_rand({hidden, head_dim}));
+    global_weights.push_back(cuda_rand({experts, hidden}));
+    global_weights.push_back(cuda_rand({1, hidden}));
+    global_weights.push_back(cuda_rand({intermediate, hidden}));
+    global_weights.push_back(cuda_rand({intermediate, hidden}));
+    global_weights.push_back(cuda_rand({hidden, intermediate}));
+    global_weights.push_back(cuda_rand({experts, 2 * intermediate, hidden}));
+    global_weights.push_back(cuda_rand({experts, hidden, intermediate}));
+    for (auto& weight : global_weights) weight.set_requires_grad(false);
+
+    assert(max_abs_diff(
+        global_weights[13].narrow(0, 0, 1),
+        global_weights[13].narrow(0, 1, 1)) > 0.0);
+    assert(max_abs_diff(
+        global_weights[14].narrow(0, 0, 1),
+        global_weights[14].narrow(0, 1, 1)) > 0.0);
+
+    std::vector<at::Tensor> distributed_weights = global_weights;
+    distributed_weights[13] = global_weights[13].narrow(0, rank, 1).contiguous();
+    distributed_weights[14] = global_weights[14].narrow(0, rank, 1).contiguous();
+    auto distributed_weight_ptrs = tensor_ptrs(distributed_weights);
+    auto reference_weight_ptrs = tensor_ptrs(global_weights);
+
+    auto embed = cuda_rand({vocab, hidden});
+    auto final_norm = cuda_rand({hidden});
+    auto lm_head = cuda_rand({vocab, hidden});
+    embed.set_requires_grad(false);
+    final_norm.set_requires_grad(false);
+    lm_head.set_requires_grad(false);
+
+    LayerConfig distributed_config{};
+    distributed_config.layer_type = 0;
+    distributed_config.num_heads = 1;
+    distributed_config.num_kv_heads = 1;
+    distributed_config.head_dim = head_dim;
+    distributed_config.partial_rotary_factor = 1.0;
+    distributed_config.rope_theta = 10000.0;
+    distributed_config.rms_eps = 1e-5;
+    distributed_config.num_experts = experts;
+    // With two experts, top_k=2 guarantees every token exercises both EP
+    // shards and both corresponding expert LoRA optimizer paths.
+    distributed_config.top_k = 2;
+    distributed_config.moe_intermediate = intermediate;
+    distributed_config.expert_start = rank;
+    distributed_config.expert_count = 1;
+    distributed_config.norm_topk_prob = 1;
+    distributed_config.nccl_comm = nullptr;
+    distributed_config.nccl_stream = nullptr;
+
+    LayerConfig reference_config = distributed_config;
+    reference_config.expert_start = 0;
+    reference_config.expert_count = experts;
+
+    const int64_t target_layer = 0;
+    const char* targets = "experts_gate_up_proj,experts_down_proj";
+    void* distributed_ctx = qwen36_create_training_context(
+        distributed_weight_ptrs.data(),
+        static_cast<int64_t>(distributed_weight_ptrs.size()),
+        &embed, &final_norm, &lm_head, &distributed_config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, targets);
+    assert(distributed_ctx);
+
+    // The reference owns all experts and is explicitly process-local. Restore
+    // the distributed EP topology before initializing its communicator below.
+    setenv("WORLD_SIZE", "1", 1);
+    setenv("RANK", "0", 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    void* reference_ctx = qwen36_create_training_context(
+        reference_weight_ptrs.data(),
+        static_cast<int64_t>(reference_weight_ptrs.size()),
+        &embed, &final_norm, &lm_head, &reference_config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, targets);
+    assert(reference_ctx);
+    const std::string distributed_rank = std::to_string(rank);
+    setenv("WORLD_SIZE", "2", 1);
+    setenv("RANK", distributed_rank.c_str(), 1);
+    setenv("EP_SIZE", "2", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    setenv("RUSTRAIN_EP_RANK", distributed_rank.c_str(), 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+
+    auto opts = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat);
+    auto global_gate_up_a =
+        ((at::arange(experts * lora_rank * hidden, opts) + 1.0) * 1e-4)
+            .reshape({experts, lora_rank, hidden}).to(at::kBFloat16);
+    auto global_gate_up_b =
+        ((at::arange(experts * 2 * intermediate * lora_rank, opts) + 3.0) * 5e-5)
+            .reshape({experts, 2 * intermediate, lora_rank}).to(at::kBFloat16);
+    auto global_down_a =
+        ((at::arange(experts * lora_rank * intermediate, opts) + 5.0) * 8e-5)
+            .reshape({experts, lora_rank, intermediate}).to(at::kBFloat16);
+    auto global_down_b =
+        ((at::arange(experts * hidden * lora_rank, opts) + 7.0) * 6e-5)
+            .reshape({experts, hidden, lora_rank}).to(at::kBFloat16);
+
+    assert(max_abs_diff(
+        global_gate_up_a.narrow(0, 0, 1),
+        global_gate_up_a.narrow(0, 1, 1)) > 0.0);
+    assert(max_abs_diff(
+        global_down_b.narrow(0, 0, 1),
+        global_down_b.narrow(0, 1, 1)) > 0.0);
+
+    auto local_gate_up_a = global_gate_up_a.narrow(0, rank, 1).contiguous();
+    auto local_gate_up_b = global_gate_up_b.narrow(0, rank, 1).contiguous();
+    auto local_down_a = global_down_a.narrow(0, rank, 1).contiguous();
+    auto local_down_b = global_down_b.narrow(0, rank, 1).contiguous();
+    set_expert_lora(distributed_ctx,
+        local_gate_up_a, local_gate_up_b, local_down_a, local_down_b);
+    set_expert_lora(reference_ctx,
+        global_gate_up_a, global_gate_up_b, global_down_a, global_down_b);
+
+    auto distributed_lora = get_expert_lora(distributed_ctx);
+    auto reference_lora = get_expert_lora(reference_ctx);
+    assert(distributed_lora.gate_up_a->sizes() ==
+        at::IntArrayRef({1, lora_rank, hidden}));
+    assert(reference_lora.gate_up_a->sizes() ==
+        at::IntArrayRef({experts, lora_rank, hidden}));
+    assert(distributed_lora.gate_up_b->sizes() ==
+        at::IntArrayRef({1, 2 * intermediate, lora_rank}));
+    assert(distributed_lora.down_a->sizes() ==
+        at::IntArrayRef({1, lora_rank, intermediate}));
+    assert(distributed_lora.down_b->sizes() ==
+        at::IntArrayRef({1, hidden, lora_rank}));
+
+    auto gate_up_a_before = distributed_lora.gate_up_a->clone();
+    auto gate_up_b_before = distributed_lora.gate_up_b->clone();
+    auto down_a_before = distributed_lora.down_a->clone();
+    auto down_b_before = distributed_lora.down_b->clone();
+
+    // Initialize NCCL only after the reference context is fully constructed.
+    // qwen36_init_nccl mutates only the distributed context's copied configs.
+    assert(qwen36_init_nccl(distributed_ctx) == 0);
+
+    auto long_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kLong);
+    auto float_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat);
+    auto bool_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kBool);
+    if (!dynamic_only) {
+    at::Tensor input_ids;
+    at::Tensor target_mask;
+    at::Tensor attention_mask;
+    at::Tensor reference_input_ids;
+    at::Tensor reference_target_mask;
+    at::Tensor reference_attention_mask;
+    if (sharded) {
+        // Deliberately unequal supervised-token counts: rank 0 contributes one
+        // response token, rank 1 contributes three. The reference evaluates
+        // the deterministic global batch on a no-NCCL full-expert context.
+        input_ids = rank == 0
+            ? at::tensor({1, 2, 3, 4}, long_opts).reshape({1, 4})
+            : at::tensor({4, 5, 6, 7}, long_opts).reshape({1, 4});
+        target_mask = rank == 0
+            ? at::tensor({0, 1, 0, 0}, float_opts).reshape({1, 4})
+            : at::tensor({0, 1, 1, 1}, float_opts).reshape({1, 4});
+        attention_mask = at::ones({1, 4}, bool_opts);
+        reference_input_ids = at::cat({
+            at::tensor({1, 2, 3, 4}, long_opts).reshape({1, 4}),
+            at::tensor({4, 5, 6, 7}, long_opts).reshape({1, 4})}, 0);
+        reference_target_mask = at::cat({
+            at::tensor({0, 1, 0, 0}, float_opts).reshape({1, 4}),
+            at::tensor({0, 1, 1, 1}, float_opts).reshape({1, 4})}, 0);
+        reference_attention_mask = at::ones({2, 4}, bool_opts);
+    } else {
+        input_ids = at::tensor({1, 2, 3}, long_opts).reshape({1, 3});
+        target_mask = at::ones({1, 3}, float_opts);
+        attention_mask = at::ones({1, 3}, bool_opts);
+        reference_input_ids = input_ids;
+        reference_target_mask = target_mask;
+        reference_attention_mask = attention_mask;
+    }
+
+    const double distributed_loss = qwen36_train_step(
+        distributed_ctx, &input_ids, &target_mask, &attention_mask);
+    const double reference_loss = qwen36_train_step(
+        reference_ctx, &reference_input_ids, &reference_target_mask,
+        &reference_attention_mask);
+    c10::cuda::device_synchronize();
+    assert(distributed_loss > 0.0 && std::isfinite(distributed_loss));
+    assert(reference_loss > 0.0 && std::isfinite(reference_loss));
+
+    const auto reference_slice = [rank](const at::Tensor& tensor) {
+        return tensor.narrow(0, rank, 1);
+    };
+    const double gate_up_a_diff = max_abs_diff(
+        *distributed_lora.gate_up_a, reference_slice(*reference_lora.gate_up_a));
+    const double gate_up_b_diff = max_abs_diff(
+        *distributed_lora.gate_up_b, reference_slice(*reference_lora.gate_up_b));
+    const double down_a_diff = max_abs_diff(
+        *distributed_lora.down_a, reference_slice(*reference_lora.down_a));
+    const double down_b_diff = max_abs_diff(
+        *distributed_lora.down_b, reference_slice(*reference_lora.down_b));
+    // qwen36_train_step returns each rank's local mean in sharded mode; the
+    // global weighted scalar is intentionally compared through parameter/state
+    // parity below rather than against the full-batch reference scalar.
+    const double loss_diff = sharded ? -1.0 :
+        std::abs(distributed_loss - reference_loss);
+
+    const double gate_up_a_update = update_norm(
+        *distributed_lora.gate_up_a, gate_up_a_before);
+    const double gate_up_b_update = update_norm(
+        *distributed_lora.gate_up_b, gate_up_b_before);
+    const double down_a_update = update_norm(
+        *distributed_lora.down_a, down_a_before);
+    const double down_b_update = update_norm(
+        *distributed_lora.down_b, down_b_before);
+
+    constexpr int64_t optimizer_slots = 18;
+    std::vector<void*> distributed_m(optimizer_slots), distributed_v(optimizer_slots);
+    std::vector<void*> reference_m(optimizer_slots), reference_v(optimizer_slots);
+    assert(qwen36_export_optimizer_state(distributed_ctx,
+        distributed_m.data(), distributed_v.data(), optimizer_slots) == optimizer_slots);
+    assert(qwen36_export_optimizer_state(reference_ctx,
+        reference_m.data(), reference_v.data(), optimizer_slots) == optimizer_slots);
+
+    double optimizer_m_diff = 0.0;
+    double optimizer_v_diff = 0.0;
+    for (int64_t state_idx = 14; state_idx < optimizer_slots; ++state_idx) {
+        auto* local_m = reinterpret_cast<at::Tensor*>(distributed_m[state_idx]);
+        auto* local_v = reinterpret_cast<at::Tensor*>(distributed_v[state_idx]);
+        auto* full_m = reinterpret_cast<at::Tensor*>(reference_m[state_idx]);
+        auto* full_v = reinterpret_cast<at::Tensor*>(reference_v[state_idx]);
+        assert(local_m && local_v && full_m && full_v);
+        optimizer_m_diff = std::max(
+            optimizer_m_diff, max_abs_diff(*local_m, reference_slice(*full_m)));
+        optimizer_v_diff = std::max(
+            optimizer_v_diff, max_abs_diff(*local_v, reference_slice(*full_v)));
+    }
+
+    const double gate_up_a_adam_diff = first_adam_step_diff(
+        *distributed_lora.gate_up_a, gate_up_a_before,
+        *reinterpret_cast<at::Tensor*>(distributed_m[14]),
+        *reinterpret_cast<at::Tensor*>(distributed_v[14]));
+    const double gate_up_b_adam_diff = first_adam_step_diff(
+        *distributed_lora.gate_up_b, gate_up_b_before,
+        *reinterpret_cast<at::Tensor*>(distributed_m[15]),
+        *reinterpret_cast<at::Tensor*>(distributed_v[15]));
+    const double down_a_adam_diff = first_adam_step_diff(
+        *distributed_lora.down_a, down_a_before,
+        *reinterpret_cast<at::Tensor*>(distributed_m[16]),
+        *reinterpret_cast<at::Tensor*>(distributed_v[16]));
+    const double down_b_adam_diff = first_adam_step_diff(
+        *distributed_lora.down_b, down_b_before,
+        *reinterpret_cast<at::Tensor*>(distributed_m[17]),
+        *reinterpret_cast<at::Tensor*>(distributed_v[17]));
+
+    std::printf(
+        "native_qwen36_ep_parity rank=%d world=%d top_k=2 a2a=%d sharded=%d "
+        "loss_compare=%s local_tokens=%d "
+        "distributed_loss=%0.8f reference_loss=%0.8f loss_diff=%0.8e "
+        "gate_up_a_diff=%0.8e gate_up_b_diff=%0.8e "
+        "down_a_diff=%0.8e down_b_diff=%0.8e "
+        "adam_m_diff=%0.8e adam_v_diff=%0.8e "
+        "adam_step_diffs=[%0.8e,%0.8e,%0.8e,%0.8e] "
+        "updates=[%0.8e,%0.8e,%0.8e,%0.8e]\n",
+        rank, world, a2a, sharded, sharded ? "skipped" : "direct",
+        sharded ? (rank == 0 ? 1 : 3) : 2,
+        distributed_loss, reference_loss, loss_diff,
+        gate_up_a_diff, gate_up_b_diff, down_a_diff, down_b_diff,
+        optimizer_m_diff, optimizer_v_diff,
+        gate_up_a_adam_diff, gate_up_b_adam_diff,
+        down_a_adam_diff, down_b_adam_diff,
+        gate_up_a_update, gate_up_b_update, down_a_update, down_b_update);
+    std::fflush(stdout);
+
+    assert(gate_up_a_update > 0.0);
+    assert(gate_up_b_update > 0.0);
+    assert(down_a_update > 0.0);
+    assert(down_b_update > 0.0);
+    assert(sharded || loss_diff <= 2e-2);
+    assert(gate_up_a_diff <= param_tol);
+    assert(gate_up_b_diff <= param_tol);
+    assert(down_a_diff <= param_tol);
+    assert(down_b_diff <= param_tol);
+    assert(optimizer_m_diff <= m_tol);
+    assert(optimizer_v_diff <= v_tol);
+    assert(gate_up_a_adam_diff <= 1e-5);
+    assert(gate_up_b_adam_diff <= 1e-5);
+    assert(down_a_adam_diff <= 1e-5);
+    assert(down_b_adam_diff <= 1e-5);
+    }
+
+    if (sharded) {
+        // Dynamic tenant rows use the source flattened token index that
+        // sharded A2A already transports. Rank 0 and rank 1 deliberately swap
+        // token counts [1,2]/[2,1] for the two adapters, exercising the
+        // all-reduced per-adapter denominator and owner-local expert update.
+        const int64_t dynamic_targets[] = {0};
+        const char* dynamic_modules =
+            "experts_gate_up_proj,experts_down_proj";
+        const int64_t adapter_a = qwen36_add_lora(
+            distributed_ctx, lora_rank, 1.0, dynamic_targets, 1,
+            dynamic_modules);
+        const int64_t adapter_b = qwen36_add_lora(
+            distributed_ctx, lora_rank, 1.0, dynamic_targets, 1,
+            dynamic_modules);
+        const int64_t reference_adapter_a = qwen36_add_lora(
+            reference_ctx, lora_rank, 1.0, dynamic_targets, 1,
+            dynamic_modules);
+        const int64_t reference_adapter_b = qwen36_add_lora(
+            reference_ctx, lora_rank, 1.0, dynamic_targets, 1,
+            dynamic_modules);
+        assert(adapter_a > 0 && adapter_b > adapter_a);
+        assert(reference_adapter_a > 0 &&
+            reference_adapter_b > reference_adapter_a);
+        auto dynamic_tensor = [&](void* ctx, int64_t adapter,
+                                  const char* module, int b) {
+            auto* ptr = qwen36_get_adapter_lora_tensor(
+                ctx, adapter, 0, module, b);
+            assert(ptr);
+            return reinterpret_cast<at::Tensor*>(ptr);
+        };
+        auto initialize_dynamic_adapter = [&](int64_t distributed_adapter,
+                                              int64_t reference_adapter,
+                                              double offset) {
+            auto set_pair = [&](const char* module, int b,
+                                std::vector<int64_t> shape) {
+                int64_t numel = 1;
+                for (int64_t dim : shape) numel *= dim;
+                auto global = ((at::arange(numel, opts) + offset) * 5e-5)
+                    .reshape(shape).to(at::kBFloat16);
+                auto local = global.narrow(0, rank, 1).contiguous();
+                assert(qwen36_set_adapter_lora_tensor(
+                    distributed_ctx, distributed_adapter, 0, module, b,
+                    &local) == 0);
+                assert(qwen36_set_adapter_lora_tensor(
+                    reference_ctx, reference_adapter, 0, module, b,
+                    &global) == 0);
+            };
+            set_pair("experts_gate_up_proj", 0,
+                {experts, lora_rank, hidden});
+            set_pair("experts_gate_up_proj", 1,
+                {experts, 2 * intermediate, lora_rank});
+            set_pair("experts_down_proj", 0,
+                {experts, lora_rank, intermediate});
+            set_pair("experts_down_proj", 1,
+                {experts, hidden, lora_rank});
+        };
+        initialize_dynamic_adapter(
+            adapter_a, reference_adapter_a, 11.0);
+        initialize_dynamic_adapter(
+            adapter_b, reference_adapter_b, 1011.0);
+        auto* dynamic_a_gate = dynamic_tensor(
+            distributed_ctx, adapter_a, "experts_gate_up_proj", 1);
+        auto* dynamic_b_gate = dynamic_tensor(
+            distributed_ctx, adapter_b, "experts_gate_up_proj", 1);
+        auto dynamic_a_before = dynamic_a_gate->clone();
+        auto dynamic_b_before = dynamic_b_gate->clone();
+
+        auto dynamic_ids = at::tensor(
+            {1, 2, 3, 4, 4, 5, 6, 7}, long_opts).reshape({2, 4});
+        auto dynamic_mask = rank == 0
+            ? at::tensor({0, 1, 0, 0, 0, 1, 1, 0}, float_opts)
+                .reshape({2, 4})
+            : at::tensor({0, 0, 1, 1, 0, 0, 0, 1}, float_opts)
+                .reshape({2, 4});
+        auto dynamic_attention = at::ones({2, 4}, bool_opts);
+        const double dynamic_loss = qwen36_train_multi_lora(
+            distributed_ctx, &dynamic_ids, &dynamic_mask,
+            &dynamic_attention, 2, static_cast<int32_t>(lora_rank));
+        auto reference_dynamic_mask = at::tensor(
+            {0, 1, 1, 1, 0, 1, 1, 1}, float_opts).reshape({2, 4});
+        const double reference_dynamic_loss = dynamic_only
+            ? qwen36_train_multi_lora(
+                reference_ctx, &dynamic_ids, &reference_dynamic_mask,
+                &dynamic_attention, 2, static_cast<int32_t>(lora_rank))
+            : -1.0;
+        c10::cuda::device_synchronize();
+        assert(dynamic_loss > 0.0 && std::isfinite(dynamic_loss));
+        assert(!dynamic_only || (reference_dynamic_loss > 0.0 &&
+            std::isfinite(reference_dynamic_loss)));
+        assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_a) == 1);
+        assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_b) == 1);
+        const double dynamic_update_a = update_norm(
+            *dynamic_a_gate, dynamic_a_before);
+        const double dynamic_update_b = update_norm(
+            *dynamic_b_gate, dynamic_b_before);
+        assert(dynamic_update_a > 0.0 && dynamic_update_b > 0.0);
+
+        double dynamic_param_diff = -1.0;
+        double dynamic_m_diff = -1.0;
+        double dynamic_v_diff = -1.0;
+        if (dynamic_only) {
+            const auto reference_slice = [rank](const at::Tensor& tensor) {
+                return tensor.narrow(0, rank, 1);
+            };
+            auto optimizer_tensor = [](void* ctx, int64_t adapter,
+                                       const char* module, int b, int is_v) {
+                auto* ptr = qwen36_get_adapter_optimizer_tensor(
+                    ctx, adapter, 0, module, b, is_v);
+                assert(ptr);
+                return reinterpret_cast<at::Tensor*>(ptr);
+            };
+            const int64_t distributed_adapters[] = {adapter_a, adapter_b};
+            const int64_t reference_adapters[] = {
+                reference_adapter_a, reference_adapter_b};
+            dynamic_param_diff = 0.0;
+            dynamic_m_diff = 0.0;
+            dynamic_v_diff = 0.0;
+            for (int adapter_index = 0; adapter_index < 2; ++adapter_index) {
+                for (const char* module : {
+                         "experts_gate_up_proj", "experts_down_proj"}) {
+                    auto* distributed_b = dynamic_tensor(
+                        distributed_ctx, distributed_adapters[adapter_index],
+                        module, 1);
+                    auto* reference_b = dynamic_tensor(
+                        reference_ctx, reference_adapters[adapter_index],
+                        module, 1);
+                    dynamic_param_diff = std::max(
+                        dynamic_param_diff,
+                        max_abs_diff(*distributed_b,
+                            reference_slice(*reference_b)));
+                    for (int is_v = 0; is_v < 2; ++is_v) {
+                        auto* distributed_state = optimizer_tensor(
+                            distributed_ctx, distributed_adapters[adapter_index],
+                            module, 1, is_v);
+                        auto* reference_state = optimizer_tensor(
+                            reference_ctx, reference_adapters[adapter_index],
+                            module, 1, is_v);
+                        double diff = max_abs_diff(
+                            *distributed_state,
+                            reference_slice(*reference_state));
+                        if (is_v) dynamic_v_diff = std::max(dynamic_v_diff, diff);
+                        else dynamic_m_diff = std::max(dynamic_m_diff, diff);
+                    }
+                }
+            }
+            std::fprintf(stderr,
+                "native_qwen36_dynamic_reference rank=%d "
+                "param_diff=%0.8e m_diff=%0.8e v_diff=%0.8e\n",
+                rank, dynamic_param_diff, dynamic_m_diff, dynamic_v_diff);
+            assert(dynamic_param_diff <= 2e-3);
+            assert(dynamic_m_diff <= 5e-3);
+            assert(dynamic_v_diff <= 1e-5);
+        }
+
+        // A tenant with no supervised tokens on any rank is a valid no-op:
+        // it must not abort the logical step or advance its private Adam
+        // clock. This also covers the final chunk when n_max is smaller than
+        // the registry size.
+        const int64_t adapter_c = qwen36_add_lora(
+            distributed_ctx, lora_rank, 1.0, dynamic_targets, 1,
+            dynamic_modules);
+        assert(adapter_c > adapter_b);
+        auto* dynamic_c_gate = dynamic_tensor(
+            distributed_ctx, adapter_c, "experts_gate_up_proj", 1);
+        auto dynamic_c_before = dynamic_c_gate->clone();
+        auto zero_ids = at::tensor({
+            1, 2, 3, 4, 4, 5, 6, 7, 1, 2, 3, 4}, long_opts)
+            .reshape({3, 4});
+        auto zero_mask = at::tensor({
+            0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0}, float_opts)
+            .reshape({3, 4});
+        auto zero_attention = at::ones({3, 4}, bool_opts);
+        const double zero_loss = qwen36_train_multi_lora(
+            distributed_ctx, &zero_ids, &zero_mask,
+            &zero_attention, 3, static_cast<int32_t>(lora_rank));
+        assert(zero_loss > 0.0 && std::isfinite(zero_loss));
+        assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_a) == 2);
+        assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_b) == 2);
+        assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_c) == 0);
+        assert(update_norm(*dynamic_c_gate, dynamic_c_before) == 0.0);
+
+        // Ordinary single-adapter entry points must not silently pick an
+        // arbitrary tenant or crash in a batched BMM after registration.
+        auto ordinary_ids = dynamic_ids.narrow(0, 0, 1).contiguous();
+        auto ordinary_mask = dynamic_mask.narrow(0, 0, 1).contiguous();
+        auto ordinary_attention = dynamic_attention.narrow(0, 0, 1).contiguous();
+        assert(qwen36_train_step(
+            distributed_ctx, &ordinary_ids, &ordinary_mask,
+            &ordinary_attention) < 0.0);
+        assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_a) == 2);
+        assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_b) == 2);
+        std::printf(
+            "native_qwen36_dynamic_sharded rank=%d world=%d loss=%0.8f "
+            "adapter_steps=[%ld,%ld,%ld] updates=[%0.8e,%0.8e] "
+            "zero_loss=%0.8f reference_loss=%0.8f "
+            "param_diff=%0.8e m_diff=%0.8e v_diff=%0.8e\n",
+            rank, world, dynamic_loss,
+            (long)qwen36_get_adapter_step_count(distributed_ctx, adapter_a),
+            (long)qwen36_get_adapter_step_count(distributed_ctx, adapter_b),
+            (long)qwen36_get_adapter_step_count(distributed_ctx, adapter_c),
+            dynamic_update_a, dynamic_update_b, zero_loss,
+            reference_dynamic_loss, dynamic_param_diff,
+            dynamic_m_diff, dynamic_v_diff);
+        std::fflush(stdout);
+    }
+
+    qwen36_free_training_context(reference_ctx);
+    qwen36_free_training_context(distributed_ctx);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_mtp_dp_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_mtp_dp_smoke.cpp
new file mode 100644
index 00000000..01e1a694
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_mtp_dp_smoke.cpp
@@ -0,0 +1,453 @@
+#include <ATen/ATen.h>
+
+#include <algorithm>
+#include <array>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <initializer_list>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_parallel_nccl(
+    void*, int32_t, int32_t, int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t, int32_t, int32_t, int32_t);
+extern "C" int32_t qwen36_set_mtp_weights(
+    void*, void*, void*, void*, void*, void**, int64_t, void*, int64_t);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" void* qwen36_get_adapter_optimizer_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t);
+extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t);
+extern "C" int32_t qwen36_train_multi_lora_selected_v3(
+    void*, void*, void*, void*, const int64_t*, int32_t,
+    double*, double*, int32_t);
+extern "C" void qwen36_free_training_context(void*);
+
+static constexpr int32_t kDataParallel = 1 << 1;
+static constexpr int64_t kHidden = 8;
+static constexpr int64_t kHeads = 4;
+static constexpr int64_t kKvHeads = 2;
+static constexpr int64_t kHeadDim = 2;
+static constexpr int64_t kIntermediate = 12;
+static constexpr int64_t kVocab = 16;
+static constexpr int64_t kLoraRank = 2;
+static constexpr double kLearningRate = 1e-3;
+static constexpr double kBeta1 = 0.9;
+static constexpr double kBeta2 = 0.999;
+static constexpr double kAdamEps = 1e-8;
+
+static at::Tensor values(std::initializer_list<int64_t> shape, double scale) {
+    int64_t count = 1;
+    for (const auto dim : shape) count *= dim;
+    return ((at::arange(count, at::TensorOptions().device(at::kCUDA)
+        .dtype(at::kFloat)).remainder(23) - 11.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) {
+    return (lhs.to(at::kFloat) - rhs.to(at::kFloat)).abs().max().item<double>();
+}
+
+struct ModelFixture {
+    std::vector<at::Tensor> weights;
+    std::vector<void*> weight_ptrs;
+    at::Tensor embed;
+    at::Tensor final_norm;
+    at::Tensor lm_head;
+    at::Tensor mtp_fc;
+    at::Tensor mtp_pre_emb;
+    at::Tensor mtp_pre_hidden;
+    at::Tensor mtp_norm;
+    LayerConfig config{};
+
+    ModelFixture()
+        : embed(values({kVocab, kHidden}, .020)),
+          final_norm(at::ones({kHidden}, at::TensorOptions().device(at::kCUDA)
+              .dtype(at::kBFloat16))),
+          lm_head(values({kVocab, kHidden}, .015)),
+          mtp_fc(values({kHidden, 2 * kHidden}, .006)),
+          mtp_pre_emb(at::ones({kHidden}, at::TensorOptions().device(at::kCUDA)
+              .dtype(at::kBFloat16))),
+          mtp_pre_hidden(at::ones({kHidden}, at::TensorOptions().device(at::kCUDA)
+              .dtype(at::kBFloat16))),
+          mtp_norm(at::ones({kHidden}, at::TensorOptions().device(at::kCUDA)
+              .dtype(at::kBFloat16))) {
+        weights.push_back(at::ones({kHidden}, at::TensorOptions()
+            .device(at::kCUDA).dtype(at::kBFloat16)));
+        weights.push_back(at::ones({kHidden}, at::TensorOptions()
+            .device(at::kCUDA).dtype(at::kBFloat16)));
+        weights.push_back(values({2 * kHeads * kHeadDim, kHidden}, .010));
+        weights.push_back(at::ones({kHeadDim}, at::TensorOptions()
+            .device(at::kCUDA).dtype(at::kBFloat16)));
+        weights.push_back(values({kKvHeads * kHeadDim, kHidden}, .012));
+        weights.push_back(at::ones({kHeadDim}, at::TensorOptions()
+            .device(at::kCUDA).dtype(at::kBFloat16)));
+        weights.push_back(values({kKvHeads * kHeadDim, kHidden}, .008));
+        weights.push_back(values({kHidden, kHeads * kHeadDim}, .011));
+        weights.push_back(values({kIntermediate, kHidden}, .009));
+        weights.push_back(values({kIntermediate, kHidden}, .007));
+        weights.push_back(values({kHidden, kIntermediate}, .010));
+        for (auto& tensor : weights) tensor.set_requires_grad(false);
+        for (auto* tensor : {&embed, &final_norm, &lm_head, &mtp_fc,
+                             &mtp_pre_emb, &mtp_pre_hidden, &mtp_norm})
+            tensor->set_requires_grad(false);
+        weight_ptrs = pointers(weights);
+
+        config.layer_type = 0;
+        config.num_heads = kHeads;
+        config.num_kv_heads = kKvHeads;
+        config.head_dim = kHeadDim;
+        config.partial_rotary_factor = 1.0;
+        config.rope_theta = 10000.0;
+        config.rms_eps = 1e-5;
+        config.intermediate_size = kIntermediate;
+    }
+};
+
+struct LoraFixture {
+    at::Tensor a;
+    at::Tensor b;
+};
+
+static LoraFixture lora_fixture(double a_scale, double b_scale) {
+    return {
+        values({kLoraRank, kHidden}, a_scale),
+        values({2 * kHeads * kHeadDim, kLoraRank}, b_scale),
+    };
+}
+
+static void install_lora(
+    void* context, int64_t adapter, const LoraFixture& fixture
+) {
+    auto a = fixture.a;
+    auto b = fixture.b;
+    assert(qwen36_set_adapter_lora_tensor(
+        context, adapter, 0, "q_proj", 0, &a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        context, adapter, 0, "q_proj", 1, &b) == 0);
+}
+
+static void set_distributed_environment(int rank, int local_rank) {
+    setenv("WORLD_SIZE", "2", 1);
+    setenv("RANK", std::to_string(rank).c_str(), 1);
+    setenv("LOCAL_RANK", std::to_string(local_rank).c_str(), 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("CP_SIZE", "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "2", 1);
+    setenv("PP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    setenv("RUSTRAIN_CP_RANK", "0", 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", std::to_string(rank).c_str(), 1);
+    setenv("RUSTRAIN_PP_RANK", "0", 1);
+    setenv("RUSTRAIN_DATA_PARALLEL", "1", 1);
+    unsetenv("QWEN36_DISABLE_MTP");
+    setenv("QWEN36_MTP_LOSS_SCALE", "0.4", 1);
+}
+
+static void set_reference_environment(int local_rank) {
+    setenv("WORLD_SIZE", "1", 1);
+    setenv("RANK", "0", 1);
+    setenv("LOCAL_RANK", std::to_string(local_rank).c_str(), 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("CP_SIZE", "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("PP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    setenv("RUSTRAIN_CP_RANK", "0", 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    setenv("RUSTRAIN_PP_RANK", "0", 1);
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    unsetenv("QWEN36_DISABLE_MTP");
+    setenv("QWEN36_MTP_LOSS_SCALE", "0.4", 1);
+}
+
+static void* create_context(ModelFixture& model, bool data_parallel) {
+    const int64_t target_layer = 0;
+    const int32_t flags = data_parallel ? kDataParallel : 0;
+    return qwen36_create_training_context_ex(
+        model.weight_ptrs.data(), model.weight_ptrs.size(), &model.embed,
+        &model.final_norm, &model.lm_head, &model.config, 1,
+        static_cast<int32_t>(at::kBFloat16), 1.0, kLearningRate,
+        kBeta1, kBeta2, kAdamEps, kVocab, 1e-5, kLoraRank,
+        &target_layer, 1, "q_proj", flags);
+}
+
+static void install_mtp(void* context, ModelFixture& model) {
+    auto mtp_weight_ptrs = pointers(model.weights);
+    assert(qwen36_set_mtp_weights(
+        context, &model.mtp_fc, &model.mtp_pre_emb,
+        &model.mtp_pre_hidden, &model.mtp_norm,
+        mtp_weight_ptrs.data(), mtp_weight_ptrs.size(),
+        &model.config, 1) == 0);
+}
+
+struct Batch {
+    at::Tensor ids;
+    at::Tensor targets;
+    at::Tensor attention;
+};
+
+static Batch tenant_one_batch() {
+    auto ids = at::tensor({1, 2, 3, 4, 5, 6},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 6});
+    auto targets = at::tensor({0., 1., 1., 1., 1., 1.},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)).reshape({1, 6});
+    return {ids, targets, at::ones({1, 6},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool))};
+}
+
+static Batch tenant_two_batch() {
+    auto ids = at::tensor({7, 8, 9, 10, 11, 12},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 6});
+    auto targets = at::tensor({0., 1., 1., 0., 0., 0.},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)).reshape({1, 6});
+    return {ids, targets, at::ones({1, 6},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool))};
+}
+
+static Batch distributed_batch(int dp_rank) {
+    auto tenant_one = tenant_one_batch();
+    auto tenant_two = tenant_two_batch();
+    auto zero_one = at::zeros_like(tenant_one.targets);
+    auto zero_two = at::zeros_like(tenant_two.targets);
+    return {
+        at::cat({tenant_one.ids, tenant_two.ids}, 0),
+        dp_rank == 0
+            ? at::cat({tenant_one.targets, zero_two}, 0)
+            : at::cat({zero_one, tenant_two.targets}, 0),
+        at::cat({tenant_one.attention, tenant_two.attention}, 0),
+    };
+}
+
+struct AdapterState {
+    std::array<at::Tensor, 2> parameters;
+    std::array<at::Tensor, 2> adam_m;
+    std::array<at::Tensor, 2> adam_v;
+};
+
+static AdapterState adapter_state(void* context, int64_t adapter) {
+    AdapterState result;
+    for (int is_b = 0; is_b < 2; ++is_b) {
+        auto* parameter = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_lora_tensor(
+                context, adapter, 0, "q_proj", is_b));
+        auto* adam_m = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_optimizer_tensor(
+                context, adapter, 0, "q_proj", is_b, 0));
+        auto* adam_v = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_optimizer_tensor(
+                context, adapter, 0, "q_proj", is_b, 1));
+        assert(parameter && adam_m && adam_v);
+        result.parameters[is_b] = parameter->clone();
+        result.adam_m[is_b] = adam_m->clone();
+        result.adam_v[is_b] = adam_v->clone();
+    }
+    return result;
+}
+
+static double state_diff(const AdapterState& lhs, const AdapterState& rhs) {
+    double result = 0.0;
+    for (int index = 0; index < 2; ++index) {
+        result = std::max(result,
+            max_diff(lhs.parameters[index], rhs.parameters[index]));
+        result = std::max(result, max_diff(lhs.adam_m[index], rhs.adam_m[index]));
+        result = std::max(result, max_diff(lhs.adam_v[index], rhs.adam_v[index]));
+    }
+    return result;
+}
+
+struct ReferenceResult {
+    void* context;
+    int64_t adapter;
+    double loss;
+    AdapterState state;
+};
+
+static ReferenceResult run_reference(
+    ModelFixture& model, const LoraFixture& lora, Batch& batch, bool enable_mtp
+) {
+    void* context = create_context(model, false);
+    assert(context);
+    if (enable_mtp) install_mtp(context, model);
+    const int64_t target_layer = 0;
+    const int64_t adapter = qwen36_add_lora(
+        context, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    assert(adapter > 0);
+    install_lora(context, adapter, lora);
+    double aggregate = -1.0;
+    double loss = -1.0;
+    assert(qwen36_train_multi_lora_selected_v3(
+        context, &batch.ids, &batch.targets, &batch.attention,
+        &adapter, 1, &aggregate, &loss, 1) == 0);
+    assert(std::isfinite(aggregate) && std::isfinite(loss));
+    assert(std::abs(aggregate - loss) < 1e-12);
+    return {context, adapter, loss, adapter_state(context, adapter)};
+}
+
+int main() {
+    const int rank = std::atoi(std::getenv("RANK"));
+    const int world = std::atoi(std::getenv("WORLD_SIZE"));
+    const int local_rank = std::atoi(
+        std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0");
+    assert(world == 2 && rank >= 0 && rank < world);
+    qwen36_set_cuda_device(local_rank);
+    set_distributed_environment(rank, local_rank);
+
+    ModelFixture model;
+    void* distributed = create_context(model, true);
+    assert(distributed);
+    install_mtp(distributed, model);
+    assert(qwen36_init_parallel_nccl(
+        distributed, rank, world,
+        0, 1, rank,
+        0, 1, rank,
+        rank, 2, 0) == 0);
+
+    const int64_t target_layer = 0;
+    const LoraFixture lora_one = lora_fixture(.0020, .0010);
+    const LoraFixture lora_two = lora_fixture(.0026, .0007);
+    const LoraFixture lora_unselected = lora_fixture(.0008, .0016);
+    const int64_t adapter_one = qwen36_add_lora(
+        distributed, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    const int64_t adapter_two = qwen36_add_lora(
+        distributed, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    const int64_t adapter_unselected = qwen36_add_lora(
+        distributed, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    assert(adapter_one > 0 && adapter_two > adapter_one &&
+        adapter_unselected > adapter_two);
+    install_lora(distributed, adapter_one, lora_one);
+    install_lora(distributed, adapter_two, lora_two);
+    install_lora(distributed, adapter_unselected, lora_unselected);
+    const auto unselected_before = adapter_state(
+        distributed, adapter_unselected);
+
+    auto local_batch = distributed_batch(rank);
+    const int64_t selected[] = {adapter_one, adapter_two};
+    double aggregate_loss = -1.0;
+    double adapter_losses[2] = {-1.0, -1.0};
+    assert(qwen36_train_multi_lora_selected_v3(
+        distributed, &local_batch.ids, &local_batch.targets,
+        &local_batch.attention, selected, 2,
+        &aggregate_loss, adapter_losses, 2) == 0);
+    assert(std::isfinite(aggregate_loss) &&
+        std::isfinite(adapter_losses[0]) &&
+        std::isfinite(adapter_losses[1]));
+
+    set_reference_environment(local_rank);
+    auto tenant_one = tenant_one_batch();
+    auto tenant_two = tenant_two_batch();
+    const auto reference_one = run_reference(
+        model, lora_one, tenant_one, true);
+    const auto reference_two = run_reference(
+        model, lora_two, tenant_two, true);
+    const auto main_only = run_reference(
+        model, lora_one, tenant_one, false);
+    set_distributed_environment(rank, local_rank);
+
+    const auto distributed_one = adapter_state(distributed, adapter_one);
+    const auto distributed_two = adapter_state(distributed, adapter_two);
+    const auto unselected_after = adapter_state(
+        distributed, adapter_unselected);
+    double parameter_diff = 0.0;
+    double adam_m_diff = 0.0;
+    double adam_v_diff = 0.0;
+    for (int index = 0; index < 2; ++index) {
+        parameter_diff = std::max({parameter_diff,
+            max_diff(distributed_one.parameters[index],
+                reference_one.state.parameters[index]),
+            max_diff(distributed_two.parameters[index],
+                reference_two.state.parameters[index])});
+        adam_m_diff = std::max({adam_m_diff,
+            max_diff(distributed_one.adam_m[index],
+                reference_one.state.adam_m[index]),
+            max_diff(distributed_two.adam_m[index],
+                reference_two.state.adam_m[index])});
+        adam_v_diff = std::max({adam_v_diff,
+            max_diff(distributed_one.adam_v[index],
+                reference_one.state.adam_v[index]),
+            max_diff(distributed_two.adam_v[index],
+                reference_two.state.adam_v[index])});
+    }
+    const double adapter_loss_diff = std::max(
+        std::abs(adapter_losses[0] - reference_one.loss),
+        std::abs(adapter_losses[1] - reference_two.loss));
+    const double aggregate_reference =
+        0.5 * (reference_one.loss + reference_two.loss);
+    const double aggregate_diff = std::abs(
+        aggregate_loss - aggregate_reference);
+    const double unselected_diff = state_diff(
+        unselected_before, unselected_after);
+    const double mtp_loss_effect = std::abs(
+        reference_one.loss - main_only.loss);
+    double mtp_parameter_effect = 0.0;
+    for (int index = 0; index < 2; ++index) {
+        mtp_parameter_effect = std::max(mtp_parameter_effect, max_diff(
+            reference_one.state.parameters[index],
+            main_only.state.parameters[index]));
+    }
+
+    std::printf(
+        "native_qwen36_mtp_dp rank=%d aggregate_diff=%0.8e "
+        "adapter_loss_diff=%0.8e parameter_diff=%0.8e m_diff=%0.8e "
+        "v_diff=%0.8e unselected_diff=%0.8e "
+        "mtp_loss_effect=%0.8e mtp_parameter_effect=%0.8e "
+        "losses=%0.8g,%0.8g\n",
+        rank, aggregate_diff, adapter_loss_diff, parameter_diff,
+        adam_m_diff, adam_v_diff, unselected_diff,
+        mtp_loss_effect, mtp_parameter_effect,
+        adapter_losses[0], adapter_losses[1]);
+    std::fflush(stdout);
+
+    assert(aggregate_diff < 1e-4 && adapter_loss_diff < 1e-4);
+    assert(parameter_diff <= 2e-3);
+    assert(adam_m_diff < 5e-5 && adam_v_diff < 5e-8);
+    assert(unselected_diff == 0.0);
+    assert(qwen36_get_adapter_step_count(distributed, adapter_one) == 1);
+    assert(qwen36_get_adapter_step_count(distributed, adapter_two) == 1);
+    assert(qwen36_get_adapter_step_count(
+        distributed, adapter_unselected) == 0);
+    assert(mtp_loss_effect > 1e-6 || mtp_parameter_effect > 0.0);
+
+    qwen36_free_training_context(main_only.context);
+    qwen36_free_training_context(reference_two.context);
+    qwen36_free_training_context(reference_one.context);
+    qwen36_free_training_context(distributed);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_mtp_dynamic_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_mtp_dynamic_smoke.cpp
new file mode 100644
index 00000000..a98c6098
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_mtp_dynamic_smoke.cpp
@@ -0,0 +1,242 @@
+#include <ATen/ATen.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <initializer_list>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" int32_t qwen36_set_mtp_weights(
+    void*, void*, void*, void*, void*, void**, int64_t, void*, int64_t);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_train_multi_lora_selected_v3(
+    void*, void*, void*, void*, const int64_t*, int32_t,
+    double*, double*, int32_t);
+extern "C" void qwen36_free_training_context(void*);
+
+static at::Tensor values(std::initializer_list<int64_t> shape, double scale) {
+    int64_t count = 1;
+    for (const auto dim : shape) count *= dim;
+    return ((at::arange(count, at::TensorOptions().device(at::kCUDA)
+        .dtype(at::kFloat)).remainder(23) - 11.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) {
+    return (lhs.to(at::kFloat) - rhs.to(at::kFloat)).abs().max().item<double>();
+}
+
+static void install_q_lora(void* ctx, int64_t adapter, const at::Tensor& a,
+                           const at::Tensor& b) {
+    auto a_copy = a;
+    auto b_copy = b;
+    assert(qwen36_set_adapter_lora_tensor(
+        ctx, adapter, 0, "q_proj", 0, &a_copy) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        ctx, adapter, 0, "q_proj", 1, &b_copy) == 0);
+}
+
+static void set_environment() {
+    setenv("WORLD_SIZE", "1", 1);
+    setenv("RANK", "0", 1);
+    setenv("LOCAL_RANK", "0", 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("CP_SIZE", "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("PP_SIZE", "1", 1);
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    unsetenv("QWEN36_DISABLE_MTP");
+    unsetenv("QWEN36_MTP_LOSS_SCALE");
+}
+
+int main() {
+    set_environment();
+    qwen36_set_cuda_device(0);
+
+    constexpr int64_t hidden = 8;
+    constexpr int64_t heads = 4;
+    constexpr int64_t kv_heads = 2;
+    constexpr int64_t head_dim = 2;
+    constexpr int64_t intermediate = 12;
+    constexpr int64_t vocab = 16;
+    constexpr int64_t rank = 2;
+
+    std::vector<at::Tensor> weights;
+    weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    weights.push_back(values({2 * heads * head_dim, hidden}, .010));
+    weights.push_back(at::ones({head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    weights.push_back(values({kv_heads * head_dim, hidden}, .012));
+    weights.push_back(at::ones({head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    weights.push_back(values({kv_heads * head_dim, hidden}, .008));
+    weights.push_back(values({hidden, heads * head_dim}, .011));
+    weights.push_back(values({intermediate, hidden}, .009));
+    weights.push_back(values({intermediate, hidden}, .007));
+    weights.push_back(values({hidden, intermediate}, .010));
+    for (auto& tensor : weights) tensor.set_requires_grad(false);
+    auto embed = values({vocab, hidden}, .020);
+    auto final_norm = at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+    auto lm_head = values({vocab, hidden}, .015);
+    embed.set_requires_grad(false);
+    final_norm.set_requires_grad(false);
+    lm_head.set_requires_grad(false);
+
+    LayerConfig config{};
+    config.layer_type = 0;
+    config.num_heads = heads;
+    config.num_kv_heads = kv_heads;
+    config.head_dim = head_dim;
+    config.partial_rotary_factor = 1.0;
+    config.rope_theta = 10000.0;
+    config.rms_eps = 1e-5;
+    config.intermediate_size = intermediate;
+    const int64_t target_layer = 0;
+    auto weight_ptrs = pointers(weights);
+
+    void* ctx = qwen36_create_training_context(
+        weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, vocab, 1e-5, rank,
+        &target_layer, 1, "q_proj");
+    assert(ctx);
+
+    auto mtp_fc = values({hidden, 2 * hidden}, .006);
+    auto mtp_pre_emb = at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+    auto mtp_pre_hidden = at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+    auto mtp_norm = at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+    mtp_fc.set_requires_grad(false);
+    mtp_pre_emb.set_requires_grad(false);
+    mtp_pre_hidden.set_requires_grad(false);
+    mtp_norm.set_requires_grad(false);
+    auto mtp_weight_ptrs = pointers(weights);
+    assert(qwen36_set_mtp_weights(
+        ctx, &mtp_fc, &mtp_pre_emb, &mtp_pre_hidden, &mtp_norm,
+        mtp_weight_ptrs.data(), mtp_weight_ptrs.size(), &config, 1) == 0);
+
+    auto q_a = values({rank, hidden}, .002);
+    auto q_b = values({2 * heads * head_dim, rank}, .001);
+    const int64_t adapter_one = qwen36_add_lora(
+        ctx, rank, 2.0, &target_layer, 1, "q_proj");
+    const int64_t adapter_two = qwen36_add_lora(
+        ctx, rank, 2.0, &target_layer, 1, "q_proj");
+    const int64_t adapter_three = qwen36_add_lora(
+        ctx, rank, 2.0, &target_layer, 1, "q_proj");
+    assert(adapter_one > 0 && adapter_two > adapter_one &&
+        adapter_three > adapter_two);
+    install_q_lora(ctx, adapter_one, q_a, q_b);
+    install_q_lora(ctx, adapter_two, q_a * 1.3, q_b * 0.7);
+    install_q_lora(ctx, adapter_three, q_a * 0.4, q_b * 1.6);
+    auto* unselected_a_before_ptr = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(ctx, adapter_three, 0, "q_proj", 0));
+    auto* unselected_b_before_ptr = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(ctx, adapter_three, 0, "q_proj", 1));
+    assert(unselected_a_before_ptr && unselected_b_before_ptr);
+    auto unselected_a_before = unselected_a_before_ptr->clone();
+    auto unselected_b_before = unselected_b_before_ptr->clone();
+
+    // Row 0 has five main-loss tokens and four MTP tokens; row 1 has two and
+    // one respectively. This makes accidental use of one shared denominator
+    // observable while keeping the first row exactly comparable to a singleton.
+    auto input_ids = at::tensor({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({2, 6});
+    auto target_mask = at::tensor({0., 1., 1., 1., 1., 1.,
+                                    0., 1., 1., 0., 0., 0.},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)).reshape({2, 6});
+    const int64_t selected_ids[] = {adapter_one, adapter_two};
+    double aggregate = -1.0;
+    double adapter_losses[2] = {-1.0, -1.0};
+    assert(qwen36_train_multi_lora_selected_v3(
+        ctx, &input_ids, &target_mask, nullptr, selected_ids, 2,
+        &aggregate, adapter_losses, 2) == 0);
+
+    void* ref = qwen36_create_training_context(
+        weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, vocab, 1e-5, rank,
+        &target_layer, 1, "q_proj");
+    assert(ref);
+    auto ref_mtp_weight_ptrs = pointers(weights);
+    assert(qwen36_set_mtp_weights(
+        ref, &mtp_fc, &mtp_pre_emb, &mtp_pre_hidden, &mtp_norm,
+        ref_mtp_weight_ptrs.data(), ref_mtp_weight_ptrs.size(), &config, 1) == 0);
+    const int64_t ref_adapter = qwen36_add_lora(
+        ref, rank, 2.0, &target_layer, 1, "q_proj");
+    assert(ref_adapter > 0);
+    install_q_lora(ref, ref_adapter, q_a, q_b);
+    auto ref_ids = input_ids.narrow(0, 0, 1).contiguous();
+    auto ref_mask = target_mask.narrow(0, 0, 1).contiguous();
+    const int64_t ref_selected[] = {ref_adapter};
+    double ref_aggregate = -1.0;
+    double ref_loss[1] = {-1.0};
+    assert(qwen36_train_multi_lora_selected_v3(
+        ref, &ref_ids, &ref_mask, nullptr, ref_selected, 1,
+        &ref_aggregate, ref_loss, 1) == 0);
+
+    auto* updated_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(ctx, adapter_one, 0, "q_proj", 0));
+    auto* updated_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(ctx, adapter_one, 0, "q_proj", 1));
+    auto* ref_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(ref, ref_adapter, 0, "q_proj", 0));
+    auto* ref_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(ref, ref_adapter, 0, "q_proj", 1));
+    auto* unselected_a_after_ptr = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(ctx, adapter_three, 0, "q_proj", 0));
+    auto* unselected_b_after_ptr = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(ctx, adapter_three, 0, "q_proj", 1));
+    assert(updated_a && updated_b && ref_a && ref_b &&
+        unselected_a_after_ptr && unselected_b_after_ptr);
+    const double parameter_diff = std::max(
+        max_diff(*updated_a, *ref_a), max_diff(*updated_b, *ref_b));
+    const double unselected_diff = std::max(
+        max_diff(*unselected_a_after_ptr, unselected_a_before),
+        max_diff(*unselected_b_after_ptr, unselected_b_before));
+    const double loss_diff = std::abs(adapter_losses[0] - ref_loss[0]);
+    std::printf(
+        "native_qwen36_mtp_dynamic_smoke aggregate=%0.8g adapter0=%0.8g "
+        "adapter1=%0.8g ref=%0.8g loss_diff=%0.8e parameter_diff=%0.8e "
+        "unselected_diff=%0.8e\n", aggregate, adapter_losses[0],
+        adapter_losses[1], ref_loss[0], loss_diff, parameter_diff,
+        unselected_diff);
+    std::fflush(stdout);
+    assert(std::isfinite(aggregate) && std::isfinite(adapter_losses[0]) &&
+        std::isfinite(adapter_losses[1]) && std::isfinite(ref_loss[0]));
+    assert(loss_diff < 2e-5);
+    assert(parameter_diff < 2e-3);
+    assert(unselected_diff == 0.0);
+
+    qwen36_free_training_context(ref);
+    qwen36_free_training_context(ctx);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_mtp_tp_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_mtp_tp_smoke.cpp
new file mode 100644
index 00000000..7d684689
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_mtp_tp_smoke.cpp
@@ -0,0 +1,793 @@
+#include <ATen/ATen.h>
+
+#include <algorithm>
+#include <array>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <initializer_list>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_nccl(void*);
+extern "C" int32_t qwen36_set_mtp_weights(
+    void*, void*, void*, void*, void*, void**, int64_t, void*, int64_t);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" int32_t qwen36_set_lora_tensor(
+    void*, int64_t, int32_t, void*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" int64_t qwen36_export_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" void* qwen36_get_adapter_optimizer_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t);
+extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t);
+extern "C" int32_t qwen36_train_multi_lora_selected_v3(
+    void*, void*, void*, void*, const int64_t*, int32_t,
+    double*, double*, int32_t);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" void qwen36_free_training_context(void*);
+
+static constexpr int32_t kBaseTpAttention = 1 << 0;
+static constexpr int32_t kVocabParallel = 1 << 2;
+static constexpr int32_t kExpertParallel = 1 << 3;
+static constexpr int32_t kBaseTpMlp = 1 << 4;
+static constexpr int64_t kTpSize = 2;
+static constexpr int64_t kHidden = 8;
+static constexpr int64_t kHeads = 4;
+static constexpr int64_t kKvHeads = 2;
+static constexpr int64_t kHeadDim = 2;
+static constexpr int64_t kIntermediate = 12;
+static constexpr int64_t kExperts = 4;
+static constexpr int64_t kTopK = 2;
+static constexpr int64_t kMoeIntermediate = 16;
+static constexpr int64_t kSharedIntermediate = 8;
+static constexpr int64_t kVocab = 16;
+static constexpr int64_t kLoraRank = 2;
+
+static bool mtp_moe_enabled() {
+    const char* value = std::getenv("QWEN36_TEST_MTP_MOE");
+    return value && std::string(value) != "0";
+}
+
+static bool mtp_ep_enabled() {
+    const char* value = std::getenv("QWEN36_TEST_MTP_EP");
+    return value && std::string(value) != "0";
+}
+
+static bool mtp_vocab_enabled() {
+    const char* value = std::getenv("QWEN36_TEST_MTP_VOCAB");
+    return value && std::string(value) != "0";
+}
+
+static at::Tensor values(std::initializer_list<int64_t> shape, double scale,
+                         int64_t offset = 0) {
+    int64_t count = 1;
+    for (const auto dim : shape) count *= dim;
+    return (((at::arange(count, at::TensorOptions().device(at::kCUDA)
+        .dtype(at::kFloat)) + offset).remainder(23) - 11.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) {
+    return (lhs.to(at::kFloat) - rhs.to(at::kFloat))
+        .abs().max().item<double>();
+}
+
+static std::vector<at::Tensor> make_layer_weights(
+    double scale, int64_t offset, bool moe
+) {
+    auto options = at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16);
+    std::vector<at::Tensor> result;
+    result.push_back(at::ones({kHidden}, options));
+    result.push_back(at::ones({kHidden}, options));
+    result.push_back(values({2 * kHeads * kHeadDim, kHidden}, scale, offset));
+    result.push_back(at::ones({kHeadDim}, options));
+    result.push_back(values({kKvHeads * kHeadDim, kHidden}, scale * 1.2, offset + 3));
+    result.push_back(at::ones({kHeadDim}, options));
+    result.push_back(values({kKvHeads * kHeadDim, kHidden}, scale * .8, offset + 5));
+    result.push_back(values({kHidden, kHeads * kHeadDim}, scale * 1.1, offset + 7));
+    if (moe) {
+        result.push_back(values({kExperts, kHidden}, scale * .9, offset + 11));
+        result.push_back(values({1, kHidden}, scale * .7, offset + 13));
+        result.push_back(values(
+            {kSharedIntermediate, kHidden}, scale * .8, offset + 17));
+        result.push_back(values(
+            {kSharedIntermediate, kHidden}, scale * .6, offset + 19));
+        result.push_back(values(
+            {kHidden, kSharedIntermediate}, scale * .9, offset + 23));
+        result.push_back(values(
+            {kExperts, 2 * kMoeIntermediate, kHidden},
+            scale * .7, offset + 29));
+        result.push_back(values(
+            {kExperts, kHidden, kMoeIntermediate},
+            scale * .8, offset + 31));
+    } else {
+        result.push_back(values(
+            {kIntermediate, kHidden}, scale * .9, offset + 11));
+        result.push_back(values(
+            {kIntermediate, kHidden}, scale * .7, offset + 13));
+        result.push_back(values(
+            {kHidden, kIntermediate}, scale, offset + 17));
+    }
+    for (auto& tensor : result) tensor.set_requires_grad(false);
+    return result;
+}
+
+static std::vector<at::Tensor> shard_layer_weights(
+    const std::vector<at::Tensor>& full, int tp_rank, bool moe,
+    int ep_rank = 0, int ep_size = 1
+) {
+    const int64_t local_heads = kHeads / kTpSize;
+    const int64_t local_kv_heads = kKvHeads / kTpSize;
+    const int64_t local_intermediate = kIntermediate / kTpSize;
+    std::vector<at::Tensor> local{
+        full[0], full[1],
+        full[2].narrow(0, tp_rank * 2 * local_heads * kHeadDim,
+            2 * local_heads * kHeadDim).contiguous(),
+        full[3],
+        full[4].narrow(0, tp_rank * local_kv_heads * kHeadDim,
+            local_kv_heads * kHeadDim).contiguous(),
+        full[5],
+        full[6].narrow(0, tp_rank * local_kv_heads * kHeadDim,
+            local_kv_heads * kHeadDim).contiguous(),
+        full[7].narrow(1, tp_rank * local_heads * kHeadDim,
+            local_heads * kHeadDim).contiguous(),
+    };
+    if (moe) {
+        const int64_t local_moe = kMoeIntermediate / kTpSize;
+        const int64_t local_shared = kSharedIntermediate / kTpSize;
+        assert(ep_size > 0 && kExperts % ep_size == 0);
+        const int64_t local_experts = kExperts / ep_size;
+        auto local_gate_up = full[13].narrow(
+            0, ep_rank * local_experts, local_experts);
+        auto expert_gate = local_gate_up.narrow(
+            1, tp_rank * local_moe, local_moe);
+        auto expert_up = local_gate_up.narrow(
+            1, kMoeIntermediate + tp_rank * local_moe, local_moe);
+        local.push_back(full[8]);
+        local.push_back(full[9]);
+        local.push_back(full[10].narrow(
+            0, tp_rank * local_shared, local_shared).contiguous());
+        local.push_back(full[11].narrow(
+            0, tp_rank * local_shared, local_shared).contiguous());
+        local.push_back(full[12].narrow(
+            1, tp_rank * local_shared, local_shared).contiguous());
+        local.push_back(at::cat({expert_gate, expert_up}, 1).contiguous());
+        local.push_back(full[14].narrow(
+            0, ep_rank * local_experts, local_experts).narrow(
+            2, tp_rank * local_moe, local_moe).contiguous());
+    } else {
+        local.push_back(full[8].narrow(
+            0, tp_rank * local_intermediate, local_intermediate).contiguous());
+        local.push_back(full[9].narrow(
+            0, tp_rank * local_intermediate, local_intermediate).contiguous());
+        local.push_back(full[10].narrow(
+            1, tp_rank * local_intermediate, local_intermediate).contiguous());
+    }
+    return local;
+}
+
+struct ModelFixture {
+    bool moe;
+    std::vector<at::Tensor> full_weights;
+    std::vector<at::Tensor> full_mtp_weights;
+    at::Tensor embed;
+    at::Tensor final_norm;
+    at::Tensor lm_head;
+    at::Tensor mtp_fc;
+    at::Tensor mtp_pre_emb;
+    at::Tensor mtp_pre_hidden;
+    at::Tensor mtp_norm;
+    LayerConfig config{};
+
+    ModelFixture()
+        : moe(mtp_moe_enabled()),
+          full_weights(make_layer_weights(.010, 0, moe)),
+          full_mtp_weights(make_layer_weights(.008, 41, moe)),
+          embed(values({kVocab, kHidden}, .020, 19)),
+          final_norm(at::ones({kHidden}, at::TensorOptions()
+              .device(at::kCUDA).dtype(at::kBFloat16))),
+          lm_head(values({kVocab, kHidden}, .015, 29)),
+          mtp_fc(values({kHidden, 2 * kHidden}, .006, 37)),
+          mtp_pre_emb(at::ones({kHidden}, at::TensorOptions()
+              .device(at::kCUDA).dtype(at::kBFloat16))),
+          mtp_pre_hidden(at::ones({kHidden}, at::TensorOptions()
+              .device(at::kCUDA).dtype(at::kBFloat16))),
+          mtp_norm(at::ones({kHidden}, at::TensorOptions()
+              .device(at::kCUDA).dtype(at::kBFloat16))) {
+        for (auto* tensor : {&embed, &final_norm, &lm_head, &mtp_fc,
+                             &mtp_pre_emb, &mtp_pre_hidden, &mtp_norm})
+            tensor->set_requires_grad(false);
+        config.layer_type = 0;
+        config.num_heads = kHeads;
+        config.num_kv_heads = kKvHeads;
+        config.head_dim = kHeadDim;
+        config.partial_rotary_factor = 1.0;
+        config.rope_theta = 10000.0;
+        config.rms_eps = 1e-5;
+        if (moe) {
+            config.num_experts = kExperts;
+            config.top_k = kTopK;
+            config.moe_intermediate = kMoeIntermediate;
+            config.expert_start = 0;
+            config.expert_count = kExperts;
+            config.norm_topk_prob = 1;
+        } else {
+            config.intermediate_size = kIntermediate;
+        }
+    }
+};
+
+struct LoraFixture {
+    at::Tensor a;
+    at::Tensor b;
+};
+
+static LoraFixture make_lora(double a_scale, double b_scale, int64_t offset) {
+    return {
+        values({kLoraRank, kHidden}, a_scale, offset),
+        values({2 * kHeads * kHeadDim, kLoraRank}, b_scale, offset + 5),
+    };
+}
+
+static void install_lora(
+    void* context, int64_t adapter, const LoraFixture& fixture,
+    int tp_rank, bool sharded
+) {
+    auto a = fixture.a;
+    auto b = sharded
+        ? fixture.b.narrow(0,
+            tp_rank * fixture.b.size(0) / kTpSize,
+            fixture.b.size(0) / kTpSize).contiguous()
+        : fixture.b;
+    assert(qwen36_set_adapter_lora_tensor(
+        context, adapter, 0, "q_proj", 0, &a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        context, adapter, 0, "q_proj", 1, &b) == 0);
+}
+
+static void set_distributed_environment(int rank, int local_rank) {
+    setenv("WORLD_SIZE", "2", 1);
+    setenv("RANK", std::to_string(rank).c_str(), 1);
+    setenv("LOCAL_RANK", std::to_string(local_rank).c_str(), 1);
+    setenv("TP_SIZE", "2", 1);
+    setenv("CP_SIZE", "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("PP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", std::to_string(rank).c_str(), 1);
+    setenv("RUSTRAIN_CP_RANK", "0", 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    setenv("RUSTRAIN_PP_RANK", "0", 1);
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    unsetenv("QWEN36_SEQUENCE_PARALLEL");
+    if (!std::getenv("QWEN36_DISABLE_MTP")) {
+        unsetenv("QWEN36_DISABLE_MTP");
+    }
+    setenv("QWEN36_MTP_LOSS_SCALE", "0.4", 1);
+}
+
+static void set_reference_environment(int local_rank) {
+    setenv("WORLD_SIZE", "1", 1);
+    setenv("RANK", "0", 1);
+    setenv("LOCAL_RANK", std::to_string(local_rank).c_str(), 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("CP_SIZE", "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("PP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    setenv("RUSTRAIN_CP_RANK", "0", 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    setenv("RUSTRAIN_PP_RANK", "0", 1);
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    unsetenv("QWEN36_SEQUENCE_PARALLEL");
+    unsetenv("QWEN36_DISABLE_MTP");
+    setenv("QWEN36_MTP_LOSS_SCALE", "0.4", 1);
+}
+
+struct AdapterState {
+    std::array<at::Tensor, 2> parameters;
+    std::array<at::Tensor, 2> adam_m;
+    std::array<at::Tensor, 2> adam_v;
+};
+
+static AdapterState adapter_state(void* context, int64_t adapter) {
+    AdapterState result;
+    for (int is_b = 0; is_b < 2; ++is_b) {
+        auto* parameter = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_lora_tensor(
+                context, adapter, 0, "q_proj", is_b));
+        auto* adam_m = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_optimizer_tensor(
+                context, adapter, 0, "q_proj", is_b, 0));
+        auto* adam_v = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_optimizer_tensor(
+                context, adapter, 0, "q_proj", is_b, 1));
+        assert(parameter && adam_m && adam_v);
+        result.parameters[is_b] = parameter->clone();
+        result.adam_m[is_b] = adam_m->clone();
+        result.adam_v[is_b] = adam_v->clone();
+    }
+    return result;
+}
+
+static double state_diff(const AdapterState& lhs, const AdapterState& rhs) {
+    double result = 0.0;
+    for (int index = 0; index < 2; ++index) {
+        result = std::max(result,
+            max_diff(lhs.parameters[index], rhs.parameters[index]));
+        result = std::max(result,
+            max_diff(lhs.adam_m[index], rhs.adam_m[index]));
+        result = std::max(result,
+            max_diff(lhs.adam_v[index], rhs.adam_v[index]));
+    }
+    return result;
+}
+
+static void install_mtp(
+    void* context, ModelFixture& model, std::vector<at::Tensor>& weights,
+    LayerConfig& config
+) {
+    auto weight_ptrs = pointers(weights);
+    assert(qwen36_set_mtp_weights(
+        context, &model.mtp_fc, &model.mtp_pre_emb,
+        &model.mtp_pre_hidden, &model.mtp_norm,
+        weight_ptrs.data(), weight_ptrs.size(), &config, 1) == 0);
+}
+
+static void set_tp_ep_environment(
+    int rank, int local_rank, int tp_rank, int ep_rank
+) {
+    setenv("WORLD_SIZE", "4", 1);
+    setenv("RANK", std::to_string(rank).c_str(), 1);
+    setenv("LOCAL_RANK", std::to_string(local_rank).c_str(), 1);
+    setenv("TP_SIZE", "2", 1);
+    setenv("CP_SIZE", "1", 1);
+    setenv("EP_SIZE", "2", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("PP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", std::to_string(tp_rank).c_str(), 1);
+    setenv("RUSTRAIN_CP_RANK", "0", 1);
+    setenv("RUSTRAIN_EP_RANK", std::to_string(ep_rank).c_str(), 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    setenv("RUSTRAIN_PP_RANK", "0", 1);
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    unsetenv("QWEN36_SEQUENCE_PARALLEL");
+    unsetenv("QWEN36_DISABLE_MTP");
+    setenv("QWEN36_MTP_LOSS_SCALE", "0.4", 1);
+    setenv("QWEN36_EP_A2A", "1", 1);
+    setenv("QWEN36_EP_A2A_SHARDED", "1", 1);
+    setenv("QWEN36_EP_A2A_PACKED", "1", 1);
+}
+
+struct Batch {
+    at::Tensor ids;
+    at::Tensor targets;
+    at::Tensor attention;
+};
+
+static Batch ep_source_batch(int ep_rank) {
+    auto long_options = at::TensorOptions().device(at::kCUDA).dtype(at::kLong);
+    auto float_options = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat);
+    auto bool_options = at::TensorOptions().device(at::kCUDA).dtype(at::kBool);
+    if (ep_rank == 0) {
+        return {
+            at::tensor({1, 2, 3, 4, 5, 6}, long_options).reshape({1, 6}),
+            at::tensor({0., 1., 1., 1., 1., 1.}, float_options)
+                .reshape({1, 6}),
+            at::ones({1, 6}, bool_options),
+        };
+    }
+    assert(ep_rank == 1);
+    return {
+        at::tensor({7, 8, 9, 10, 11, 12}, long_options).reshape({1, 6}),
+        at::tensor({0., 1., 1., 0., 0., 0.}, float_options)
+            .reshape({1, 6}),
+        at::ones({1, 6}, bool_options),
+    };
+}
+
+static Batch full_ep_batch() {
+    auto first = ep_source_batch(0);
+    auto second = ep_source_batch(1);
+    return {
+        at::cat({first.ids, second.ids}, 0),
+        at::cat({first.targets, second.targets}, 0),
+        at::cat({first.attention, second.attention}, 0),
+    };
+}
+
+static int run_tp_ep(int rank, int world, int local_rank) {
+    assert(world == 4 && rank >= 0 && rank < world);
+    const int tp_rank = rank % 2;
+    const int ep_rank = rank / 2;
+    set_tp_ep_environment(rank, local_rank, tp_rank, ep_rank);
+
+    ModelFixture model;
+    assert(model.moe);
+    auto local_weights = shard_layer_weights(
+        model.full_weights, tp_rank, true, ep_rank, 2);
+    auto local_mtp_weights = shard_layer_weights(
+        model.full_mtp_weights, tp_rank, true, ep_rank, 2);
+    auto local_weight_ptrs = pointers(local_weights);
+    LayerConfig distributed_config = model.config;
+    distributed_config.expert_start = ep_rank * (kExperts / 2);
+    distributed_config.expert_count = kExperts / 2;
+    const int64_t target_layer = 0;
+    void* distributed = qwen36_create_training_context_ex(
+        local_weight_ptrs.data(), local_weight_ptrs.size(),
+        &model.embed, &model.final_norm, &model.lm_head,
+        &distributed_config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, kVocab, 1e-5, kLoraRank,
+        &target_layer, 1, "q_proj",
+        kBaseTpAttention | kExpertParallel | kBaseTpMlp);
+    assert(distributed);
+    install_mtp(
+        distributed, model, local_mtp_weights, distributed_config);
+    assert(qwen36_init_nccl(distributed) == 0);
+
+    const auto selected_lora = make_lora(.0020, .0010, 0);
+    const auto isolated_lora = make_lora(.0008, .0016, 31);
+    const int64_t selected_adapter = qwen36_add_lora(
+        distributed, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    const int64_t isolated_adapter = qwen36_add_lora(
+        distributed, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    assert(selected_adapter > 0 && isolated_adapter > selected_adapter);
+    install_lora(
+        distributed, selected_adapter, selected_lora, tp_rank, true);
+    install_lora(
+        distributed, isolated_adapter, isolated_lora, tp_rank, true);
+    const auto isolated_before = adapter_state(distributed, isolated_adapter);
+
+    auto local_batch = ep_source_batch(ep_rank);
+    double distributed_aggregate = -1.0;
+    double distributed_loss = -1.0;
+    assert(qwen36_train_multi_lora_selected_v3(
+        distributed, &local_batch.ids, &local_batch.targets,
+        &local_batch.attention, &selected_adapter, 1,
+        &distributed_aggregate, &distributed_loss, 1) == 0);
+    assert(std::isfinite(distributed_loss) &&
+        std::abs(distributed_aggregate - distributed_loss) < 1e-12);
+
+    set_reference_environment(local_rank);
+    auto full_weight_ptrs = pointers(model.full_weights);
+    void* reference = qwen36_create_training_context(
+        full_weight_ptrs.data(), full_weight_ptrs.size(),
+        &model.embed, &model.final_norm, &model.lm_head,
+        &model.config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, kVocab, 1e-5, kLoraRank,
+        &target_layer, 1, "q_proj");
+    assert(reference);
+    install_mtp(reference, model, model.full_mtp_weights, model.config);
+    auto reference_a_initial = selected_lora.a.clone();
+    auto reference_b_initial = selected_lora.b.clone();
+    assert(qwen36_set_lora_tensor(
+        reference, 0, 0, &reference_a_initial) == 0);
+    assert(qwen36_set_lora_tensor(
+        reference, 0, 1, &reference_b_initial) == 0);
+    auto global_batch = full_ep_batch();
+    const double reference_loss = qwen36_train_step(
+        reference, &global_batch.ids, &global_batch.targets,
+        &global_batch.attention);
+    assert(std::isfinite(reference_loss) && reference_loss > 0.0);
+
+    void* main_only = qwen36_create_training_context(
+        full_weight_ptrs.data(), full_weight_ptrs.size(),
+        &model.embed, &model.final_norm, &model.lm_head,
+        &model.config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, kVocab, 1e-5, kLoraRank,
+        &target_layer, 1, "q_proj");
+    assert(main_only);
+    auto main_only_a_initial = selected_lora.a.clone();
+    auto main_only_b_initial = selected_lora.b.clone();
+    assert(qwen36_set_lora_tensor(
+        main_only, 0, 0, &main_only_a_initial) == 0);
+    assert(qwen36_set_lora_tensor(
+        main_only, 0, 1, &main_only_b_initial) == 0);
+    const double main_only_loss = qwen36_train_step(
+        main_only, &global_batch.ids, &global_batch.targets,
+        &global_batch.attention);
+    assert(std::isfinite(main_only_loss) && main_only_loss > 0.0);
+
+    const auto selected_state = adapter_state(distributed, selected_adapter);
+    const auto isolated_after = adapter_state(distributed, isolated_adapter);
+    auto* reference_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_a(reference, 0));
+    auto* reference_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_b(reference, 0));
+    auto* main_only_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_a(main_only, 0));
+    auto* main_only_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_b(main_only, 0));
+    assert(reference_a && reference_b && main_only_a && main_only_b);
+    std::array<void*, 2> reference_m{};
+    std::array<void*, 2> reference_v{};
+    assert(qwen36_export_optimizer_state(
+        reference, reference_m.data(), reference_v.data(), 2) == 2);
+
+    const int64_t local_q_rows = 2 * (kHeads / 2) * kHeadDim;
+    auto local_reference_b = reference_b->narrow(
+        0, tp_rank * local_q_rows, local_q_rows);
+    auto local_reference_m_b = reinterpret_cast<at::Tensor*>(
+        reference_m[1])->narrow(0, tp_rank * local_q_rows, local_q_rows);
+    auto local_reference_v_b = reinterpret_cast<at::Tensor*>(
+        reference_v[1])->narrow(0, tp_rank * local_q_rows, local_q_rows);
+    const double parameter_diff = std::max(
+        max_diff(selected_state.parameters[0], *reference_a),
+        max_diff(selected_state.parameters[1], local_reference_b));
+    const double adam_m_diff = std::max(
+        max_diff(selected_state.adam_m[0],
+            *reinterpret_cast<at::Tensor*>(reference_m[0])),
+        max_diff(selected_state.adam_m[1], local_reference_m_b));
+    const double adam_v_diff = std::max(
+        max_diff(selected_state.adam_v[0],
+            *reinterpret_cast<at::Tensor*>(reference_v[0])),
+        max_diff(selected_state.adam_v[1], local_reference_v_b));
+    const double isolated_diff = state_diff(isolated_before, isolated_after);
+    const double loss_diff = std::abs(distributed_loss - reference_loss);
+    const double mtp_loss_effect = std::abs(reference_loss - main_only_loss);
+    const double mtp_parameter_effect = std::max(
+        max_diff(*reference_a, *main_only_a),
+        max_diff(*reference_b, *main_only_b));
+
+    std::printf(
+        "native_qwen36_mtp_tp_ep rank=%d tp=%d ep=%d "
+        "main_tokens=7 mtp_tokens=5 loss_diff=%0.8e "
+        "parameter_diff=%0.8e m_diff=%0.8e v_diff=%0.8e "
+        "isolated_diff=%0.8e mtp_loss_effect=%0.8e "
+        "mtp_parameter_effect=%0.8e distributed=%0.8g reference=%0.8g main_only=%0.8g\n",
+        rank, tp_rank, ep_rank, loss_diff, parameter_diff,
+        adam_m_diff, adam_v_diff, isolated_diff, mtp_loss_effect,
+        mtp_parameter_effect, distributed_loss, reference_loss, main_only_loss);
+    std::fflush(stdout);
+
+    assert(loss_diff < 1e-2);
+    assert(parameter_diff <= 3e-3);
+    assert(adam_m_diff <= 7e-3 && adam_v_diff <= 2e-5);
+    assert(isolated_diff == 0.0);
+    assert(qwen36_get_adapter_step_count(
+        distributed, selected_adapter) == 1);
+    assert(qwen36_get_adapter_step_count(
+        distributed, isolated_adapter) == 0);
+    assert(mtp_loss_effect > 1e-6 || mtp_parameter_effect > 0.0);
+
+    qwen36_free_training_context(main_only);
+    qwen36_free_training_context(reference);
+    qwen36_free_training_context(distributed);
+    return 0;
+}
+
+int main() {
+    const int rank = std::atoi(std::getenv("RANK"));
+    const int world = std::atoi(std::getenv("WORLD_SIZE"));
+    const int local_rank = std::atoi(
+        std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0");
+    qwen36_set_cuda_device(local_rank);
+    if (mtp_ep_enabled()) return run_tp_ep(rank, world, local_rank);
+    assert(world == kTpSize && rank >= 0 && rank < world);
+    set_distributed_environment(rank, local_rank);
+
+    ModelFixture model;
+    const bool vocab_parallel = mtp_vocab_enabled();
+    at::Tensor local_embed = vocab_parallel
+        ? model.embed.narrow(0, rank * (kVocab / kTpSize), kVocab / kTpSize)
+            .contiguous()
+        : model.embed;
+    at::Tensor local_lm_head = vocab_parallel
+        ? model.lm_head.narrow(0, rank * (kVocab / kTpSize), kVocab / kTpSize)
+            .contiguous()
+        : model.lm_head;
+    auto local_weights = shard_layer_weights(
+        model.full_weights, rank, model.moe);
+    auto local_mtp_weights = shard_layer_weights(
+        model.full_mtp_weights, rank, model.moe);
+    auto local_weight_ptrs = pointers(local_weights);
+    const int64_t target_layer = 0;
+    void* distributed = qwen36_create_training_context_ex(
+        local_weight_ptrs.data(), local_weight_ptrs.size(),
+        &local_embed, &model.final_norm, &local_lm_head,
+        &model.config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, kVocab, 1e-5, kLoraRank,
+        &target_layer, 1, "q_proj",
+        kBaseTpAttention | kBaseTpMlp |
+            (vocab_parallel ? kVocabParallel : 0));
+    assert(distributed);
+    install_mtp(distributed, model, local_mtp_weights, model.config);
+    assert(qwen36_init_nccl(distributed) == 0);
+
+    const auto lora_one = make_lora(.0020, .0010, 0);
+    const auto lora_two = make_lora(.0026, .0007, 17);
+    const auto lora_unselected = make_lora(.0008, .0016, 31);
+    const int64_t adapter_one = qwen36_add_lora(
+        distributed, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    const int64_t adapter_two = qwen36_add_lora(
+        distributed, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    const int64_t adapter_unselected = qwen36_add_lora(
+        distributed, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    assert(adapter_one > 0 && adapter_two > adapter_one &&
+        adapter_unselected > adapter_two);
+    install_lora(distributed, adapter_one, lora_one, rank, true);
+    install_lora(distributed, adapter_two, lora_two, rank, true);
+    install_lora(distributed, adapter_unselected, lora_unselected, rank, true);
+    const auto unselected_before = adapter_state(distributed, adapter_unselected);
+
+    auto input_ids = at::tensor({1, 2, 3, 4, 5, 6,
+                                 7, 8, 9, 10, 11, 12},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({2, 6});
+    auto target_mask = at::tensor({0., 1., 1., 1., 1., 1.,
+                                   0., 1., 1., 0., 0., 0.},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)).reshape({2, 6});
+    auto attention_mask = at::ones({2, 6},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+    const int64_t selected[] = {adapter_one, adapter_two};
+    double aggregate = -1.0;
+    double losses[2] = {-1.0, -1.0};
+    assert(qwen36_train_multi_lora_selected_v3(
+        distributed, &input_ids, &target_mask, &attention_mask,
+        selected, 2, &aggregate, losses, 2) == 0);
+
+    set_reference_environment(local_rank);
+    auto full_weight_ptrs = pointers(model.full_weights);
+    void* reference = qwen36_create_training_context(
+        full_weight_ptrs.data(), full_weight_ptrs.size(),
+        &model.embed, &model.final_norm, &model.lm_head,
+        &model.config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, kVocab, 1e-5, kLoraRank,
+        &target_layer, 1, "q_proj");
+    assert(reference);
+    install_mtp(reference, model, model.full_mtp_weights, model.config);
+    const int64_t reference_one = qwen36_add_lora(
+        reference, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    const int64_t reference_two = qwen36_add_lora(
+        reference, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    const int64_t reference_unselected = qwen36_add_lora(
+        reference, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    assert(reference_one > 0 && reference_two > reference_one &&
+        reference_unselected > reference_two);
+    install_lora(reference, reference_one, lora_one, 0, false);
+    install_lora(reference, reference_two, lora_two, 0, false);
+    install_lora(reference, reference_unselected, lora_unselected, 0, false);
+    const int64_t reference_selected[] = {reference_one, reference_two};
+    double reference_aggregate = -1.0;
+    double reference_losses[2] = {-1.0, -1.0};
+    assert(qwen36_train_multi_lora_selected_v3(
+        reference, &input_ids, &target_mask, &attention_mask,
+        reference_selected, 2, &reference_aggregate,
+        reference_losses, 2) == 0);
+
+    // A main-loss-only singleton makes a silently skipped MTP branch visible.
+    void* main_only = qwen36_create_training_context(
+        full_weight_ptrs.data(), full_weight_ptrs.size(),
+        &model.embed, &model.final_norm, &model.lm_head,
+        &model.config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, kVocab, 1e-5, kLoraRank,
+        &target_layer, 1, "q_proj");
+    assert(main_only);
+    const int64_t main_only_adapter = qwen36_add_lora(
+        main_only, kLoraRank, 2.0, &target_layer, 1, "q_proj");
+    assert(main_only_adapter > 0);
+    install_lora(main_only, main_only_adapter, lora_one, 0, false);
+    auto main_only_ids = input_ids.narrow(0, 0, 1).contiguous();
+    auto main_only_targets = target_mask.narrow(0, 0, 1).contiguous();
+    auto main_only_attention = attention_mask.narrow(0, 0, 1).contiguous();
+    double main_only_aggregate = -1.0;
+    double main_only_loss = -1.0;
+    assert(qwen36_train_multi_lora_selected_v3(
+        main_only, &main_only_ids, &main_only_targets, &main_only_attention,
+        &main_only_adapter, 1, &main_only_aggregate,
+        &main_only_loss, 1) == 0);
+
+    const auto distributed_one = adapter_state(distributed, adapter_one);
+    const auto distributed_two = adapter_state(distributed, adapter_two);
+    const auto reference_one_state = adapter_state(reference, reference_one);
+    const auto reference_two_state = adapter_state(reference, reference_two);
+    const auto main_only_state = adapter_state(main_only, main_only_adapter);
+    const auto unselected_after = adapter_state(distributed, adapter_unselected);
+    double parameter_diff = 0.0;
+    double adam_m_diff = 0.0;
+    double adam_v_diff = 0.0;
+    const int64_t local_q_rows = 2 * (kHeads / kTpSize) * kHeadDim;
+    auto compare_state = [&](const AdapterState& actual,
+                             const AdapterState& expected) {
+        for (int factor = 0; factor < 2; ++factor) {
+            auto expected_parameter = factor == 0
+                ? expected.parameters[factor]
+                : expected.parameters[factor].narrow(
+                    0, rank * local_q_rows, local_q_rows);
+            auto expected_m = factor == 0
+                ? expected.adam_m[factor]
+                : expected.adam_m[factor].narrow(
+                    0, rank * local_q_rows, local_q_rows);
+            auto expected_v = factor == 0
+                ? expected.adam_v[factor]
+                : expected.adam_v[factor].narrow(
+                    0, rank * local_q_rows, local_q_rows);
+            parameter_diff = std::max(parameter_diff,
+                max_diff(actual.parameters[factor], expected_parameter));
+            adam_m_diff = std::max(adam_m_diff,
+                max_diff(actual.adam_m[factor], expected_m));
+            adam_v_diff = std::max(adam_v_diff,
+                max_diff(actual.adam_v[factor], expected_v));
+        }
+    };
+    compare_state(distributed_one, reference_one_state);
+    compare_state(distributed_two, reference_two_state);
+
+    const double aggregate_diff = std::abs(aggregate - reference_aggregate);
+    const double loss_diff = std::max(
+        std::abs(losses[0] - reference_losses[0]),
+        std::abs(losses[1] - reference_losses[1]));
+    const double unselected_diff = state_diff(
+        unselected_before, unselected_after);
+    const double mtp_loss_effect = std::abs(
+        reference_losses[0] - main_only_loss);
+    double mtp_parameter_effect = 0.0;
+    for (int factor = 0; factor < 2; ++factor) {
+        mtp_parameter_effect = std::max(mtp_parameter_effect, max_diff(
+            reference_one_state.parameters[factor],
+            main_only_state.parameters[factor]));
+    }
+    std::printf(
+        "native_qwen36_mtp_tp mode=%s rank=%d aggregate_diff=%0.8e "
+        "loss_diff=%0.8e parameter_diff=%0.8e m_diff=%0.8e "
+        "v_diff=%0.8e unselected_diff=%0.8e mtp_loss_effect=%0.8e "
+        "mtp_parameter_effect=%0.8e losses=%0.8g,%0.8g\n",
+        vocab_parallel ? "vocab" : (model.moe ? "moe" : "dense"), rank,
+        aggregate_diff, loss_diff, parameter_diff, adam_m_diff,
+        adam_v_diff, unselected_diff, mtp_loss_effect,
+        mtp_parameter_effect, losses[0], losses[1]);
+    std::fflush(stdout);
+
+    assert(std::isfinite(aggregate) && std::isfinite(losses[0]) &&
+        std::isfinite(losses[1]));
+    assert(aggregate_diff < 5e-3 && loss_diff < 5e-3);
+    assert(parameter_diff <= (vocab_parallel ? 2.1e-3 : 2e-3));
+    assert(adam_m_diff < 5e-4 && adam_v_diff < 5e-7);
+    assert(unselected_diff == 0.0);
+    assert(qwen36_get_adapter_step_count(distributed, adapter_one) == 1);
+    assert(qwen36_get_adapter_step_count(distributed, adapter_two) == 1);
+    assert(qwen36_get_adapter_step_count(distributed, adapter_unselected) == 0);
+    assert(mtp_loss_effect > 1e-6 || mtp_parameter_effect > 0.0);
+
+    qwen36_free_training_context(main_only);
+    qwen36_free_training_context(reference);
+    qwen36_free_training_context(distributed);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_pp_cp_comm_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_pp_cp_comm_smoke.cpp
new file mode 100644
index 00000000..d7b61460
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_pp_cp_comm_smoke.cpp
@@ -0,0 +1,259 @@
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" int64_t qwen36_kernel_abi_version();
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_v2(
+    void**, int64_t, void*, void*, void*, void*, int64_t,
+    int64_t, int64_t, int32_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_parallel_nccl_v2(
+    void*, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" int32_t qwen36_attach_parallel_nccl_no_sync_v2(
+    void*, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" double qwen36_parallel_max_double(void*, double);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" double qwen36_eval_step(void*, void*, void*, void*);
+extern "C" int64_t qwen36_get_lora_count(void*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int32_t qwen36_set_adapter_id(void*, int64_t, int64_t);
+extern "C" int32_t qwen36_remove_lora(void*, int64_t);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" void qwen36_free_training_context(void*);
+
+static void set_rank_environment(int rank, int cp_rank, int pp_rank) {
+    const auto rank_string = std::to_string(rank);
+    const auto cp_string = std::to_string(cp_rank);
+    const auto pp_string = std::to_string(pp_rank);
+    setenv("RANK", rank_string.c_str(), 1);
+    setenv("WORLD_SIZE", "4", 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("CP_SIZE", "2", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("PP_SIZE", "2", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    setenv("RUSTRAIN_CP_RANK", cp_string.c_str(), 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    setenv("RUSTRAIN_PP_RANK", pp_string.c_str(), 1);
+}
+
+static void* create_empty_context(at::Tensor& embed, at::Tensor& norm,
+                                  at::Tensor& lm_head) {
+    return qwen36_create_training_context(
+        nullptr, 0, &embed, &norm, &lm_head, nullptr, 0,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, 4, 1e-6, 1,
+        nullptr, 0, "");
+}
+
+static void* create_stage_context(
+    int pp_rank, at::Tensor& embed, at::Tensor& norm, at::Tensor& lm_head,
+    std::vector<at::Tensor>& weights, LayerConfig& config
+) {
+    std::vector<void*> weight_ptrs;
+    for (auto& weight : weights) weight_ptrs.push_back(&weight);
+    const int64_t target_layer = 0;
+    return qwen36_create_training_context_v2(
+        weight_ptrs.data(), static_cast<int64_t>(weight_ptrs.size()),
+        pp_rank == 0 ? &embed : nullptr,
+        pp_rank == 1 ? &norm : nullptr,
+        pp_rank == 1 ? &lm_head : nullptr,
+        &config, 1, pp_rank, 2, pp_rank == 0 ? 1 : 2,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, 4, 1e-6, 1,
+        &target_layer, 1, "q_proj", 0);
+}
+
+static int32_t initialize_five_axis_context(
+    void* ctx, int rank, int cp_rank, int pp_rank, bool attach
+) {
+    const int cp_color = pp_rank * 2;  // [0,1] and [2,3]
+    const int pp_color = cp_rank;      // [0,2] and [1,3]
+    auto init = attach ? qwen36_attach_parallel_nccl_no_sync_v2
+                       : qwen36_init_parallel_nccl_v2;
+    return init(
+        ctx, rank, 4,
+        0, 1, 0,
+        cp_rank, 2, cp_color,
+        0, 1, 0,
+        0, 1, 0,
+        pp_rank, 2, pp_color);
+}
+
+int main() {
+    assert(qwen36_kernel_abi_version() == 31);
+    const int rank = std::atoi(std::getenv("RANK"));
+    const int local_rank = std::atoi(std::getenv("LOCAL_RANK"));
+    assert(rank >= 0 && rank < 4);
+    const int cp_rank = rank % 2;
+    const int pp_rank = rank / 2;
+    set_rank_environment(rank, cp_rank, pp_rank);
+
+    c10::cuda::CUDAGuard guard(local_rank);
+    auto options = at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16);
+    auto embed = at::zeros({4, 4}, options);
+    auto norm = at::ones({4}, options);
+    auto lm_head = at::zeros({4, 4}, options);
+
+    std::vector<at::Tensor> stage_weights = {
+        at::ones({4}, options), at::ones({4}, options),
+        at::zeros({8, 4}, options), at::ones({4}, options),
+        at::zeros({4, 4}, options), at::ones({4}, options),
+        at::zeros({4, 4}, options), at::zeros({4, 4}, options),
+        at::zeros({8, 4}, options), at::zeros({8, 4}, options),
+        at::zeros({4, 8}, options),
+    };
+    LayerConfig stage_config{};
+    stage_config.layer_type = 0;
+    stage_config.num_heads = 1;
+    stage_config.num_kv_heads = 1;
+    stage_config.head_dim = 4;
+    stage_config.partial_rotary_factor = 1.0;
+    stage_config.rope_theta = 10000.0;
+    stage_config.rms_eps = 1e-6;
+    stage_config.intermediate_size = 8;
+
+    void* premature_shadow = create_empty_context(embed, norm, lm_head);
+    assert(premature_shadow);
+    assert(initialize_five_axis_context(
+        premature_shadow, rank, cp_rank, pp_rank, true) == -1);
+    assert(std::isnan(qwen36_parallel_max_double(premature_shadow, rank)));
+    qwen36_free_training_context(premature_shadow);
+
+    const char* configured_run_id = std::getenv("RUSTRAIN_NCCL_RUN_ID");
+    const bool had_configured_run_id = configured_run_id != nullptr;
+    const std::string saved_run_id = configured_run_id ? configured_run_id : "";
+    setenv("RUSTRAIN_NCCL_RUN_ID", "..", 1);
+    void* invalid_run_id = create_empty_context(embed, norm, lm_head);
+    assert(invalid_run_id);
+    assert(initialize_five_axis_context(
+        invalid_run_id, rank, cp_rank, pp_rank, false) == -1);
+    assert(std::isnan(qwen36_parallel_max_double(invalid_run_id, rank)));
+    qwen36_free_training_context(invalid_run_id);
+    if (had_configured_run_id) {
+        setenv("RUSTRAIN_NCCL_RUN_ID", saved_run_id.c_str(), 1);
+    } else {
+        unsetenv("RUSTRAIN_NCCL_RUN_ID");
+    }
+
+    void* context = create_empty_context(embed, norm, lm_head);
+    assert(context);
+    assert(initialize_five_axis_context(
+        context, rank, cp_rank, pp_rank, false) == 0);
+    const double maximum = qwen36_parallel_max_double(context, rank);
+    assert(std::isfinite(maximum));
+    assert(maximum == 3.0);
+    assert(qwen36_train_step(context, nullptr, nullptr, nullptr) < 0.0);
+    assert(qwen36_eval_step(context, nullptr, nullptr, nullptr) < 0.0);
+
+    void* stage_context = create_stage_context(
+        pp_rank, embed, norm, lm_head, stage_weights, stage_config);
+    assert(stage_context);
+    assert(initialize_five_axis_context(
+        stage_context, rank, cp_rank, pp_rank, true) == 0);
+    assert(qwen36_get_lora_count(stage_context) == 7);
+    auto* stage_lora_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_a(stage_context, 0));
+    assert(stage_lora_a);
+    if (pp_rank == 0) {
+        assert(stage_lora_a->sizes() == at::IntArrayRef({1, 4}));
+    } else {
+        assert(stage_lora_a->dim() == 0);
+    }
+    // A single CP replica presenting a different canonical request must make
+    // every rank fail without letting unaffected PP stages skip the request
+    // consensus and deadlock.
+    const int64_t mismatched_target_layer = rank == 0 ? 1 : 0;
+    assert(qwen36_add_lora(
+        stage_context, 4, 8.0, &mismatched_target_layer, 1, "q_proj") < 0);
+    const int64_t dynamic_target_layer = 0;
+    const int64_t adapter_id = qwen36_add_lora(
+        stage_context, 4, 8.0, &dynamic_target_layer, 1, "q_proj");
+    assert(adapter_id == 1);
+    auto* dynamic_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            stage_context, adapter_id, dynamic_target_layer, "q_proj", 0));
+    if (pp_rank == 0) {
+        assert(dynamic_a);
+        assert(dynamic_a->sizes() == at::IntArrayRef({4, 4}));
+    } else {
+        assert(dynamic_a == nullptr);
+    }
+    assert(qwen36_get_adapter_lora_tensor(
+        stage_context, adapter_id, 1, "q_proj", 0) == nullptr);
+    assert(qwen36_set_adapter_id(stage_context, adapter_id, 7) == 0);
+    assert(qwen36_get_adapter_lora_tensor(
+        stage_context, adapter_id, dynamic_target_layer, "q_proj", 0) ==
+        nullptr);
+    auto* renamed_dynamic_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            stage_context, 7, dynamic_target_layer, "q_proj", 0));
+    if (pp_rank == 0) {
+        assert(renamed_dynamic_a);
+        assert(renamed_dynamic_a->sizes() == at::IntArrayRef({4, 4}));
+    } else {
+        assert(renamed_dynamic_a == nullptr);
+    }
+    assert(qwen36_remove_lora(stage_context, 7) == 1);
+    assert(qwen36_get_adapter_lora_tensor(
+        stage_context, 7, dynamic_target_layer, "q_proj", 0) == nullptr);
+
+    void* shadow = create_empty_context(embed, norm, lm_head);
+    assert(shadow);
+    assert(initialize_five_axis_context(
+        shadow, rank, cp_rank, pp_rank, true) == 0);
+    assert(qwen36_parallel_max_double(shadow, rank) == 3.0);
+
+    set_rank_environment(rank, (cp_rank + 1) % 2, pp_rank);
+    void* invalid = create_empty_context(embed, norm, lm_head);
+    assert(invalid);
+    assert(initialize_five_axis_context(
+        invalid, rank, cp_rank, pp_rank, true) == -1);
+    assert(std::isnan(qwen36_parallel_max_double(invalid, rank)));
+
+    std::printf(
+        "native_qwen36_pp_cp_comm rank=%d cp=%d pp=%d max=%0.1f ok\n",
+        rank, cp_rank, pp_rank, maximum);
+    qwen36_free_training_context(invalid);
+    qwen36_free_training_context(shadow);
+    qwen36_free_training_context(stage_context);
+    qwen36_free_training_context(context);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_pp_train_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_pp_train_smoke.cpp
new file mode 100644
index 00000000..4409c015
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_pp_train_smoke.cpp
@@ -0,0 +1,849 @@
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <utility>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" int64_t qwen36_kernel_abi_version();
+extern "C" void* qwen36_create_training_context_v2(
+    void**, int64_t, void*, void*, void*, void*, int64_t,
+    int64_t, int64_t, int32_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_parallel_nccl_v2(
+    void*, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" int32_t qwen36_attach_parallel_nccl_no_sync_v2(
+    void*, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" double qwen36_train_micro_step(
+    void*, void*, void*, void*, double, int32_t);
+extern "C" int32_t qwen36_set_pad_token_id(void*, int64_t);
+struct Qwen36PipelineWindowV1 {
+    uint32_t struct_size;
+    uint32_t version;
+    int64_t window_id;
+    int64_t num_microbatches;
+    int32_t schedule;
+    int32_t num_chunks;
+    int32_t flags;
+};
+struct Qwen36PipelineTickV1 {
+    uint32_t struct_size;
+    uint32_t version;
+    int64_t window_id;
+    int64_t forward_mb;
+    int64_t backward_mb;
+    int32_t chunk_id;
+    int32_t phase;
+    void* input_ids;
+    void* target_mask;
+    void* attention_mask;
+    double gradient_scale;
+};
+struct Qwen36PipelineResultV1 {
+    uint32_t struct_size;
+    uint32_t version;
+    int32_t status;
+    int64_t completed_fwd;
+    int64_t completed_bwd;
+    int64_t in_flight;
+    int64_t optimizer_step;
+    double loss;
+};
+extern "C" int32_t qwen36_pipeline_begin_v1(
+    void*, const Qwen36PipelineWindowV1*);
+extern "C" int32_t qwen36_pipeline_begin_dynamic_selected_v1(
+    void*, const Qwen36PipelineWindowV1*, const int64_t*, int32_t);
+extern "C" int32_t qwen36_pipeline_tick_v1(
+    void*, const Qwen36PipelineTickV1*, Qwen36PipelineResultV1*);
+extern "C" int32_t qwen36_pipeline_finish_v1(
+    void*, int32_t, Qwen36PipelineResultV1*);
+extern "C" int32_t qwen36_pipeline_finish_dynamic_report_v1(
+    void*, int32_t, Qwen36PipelineResultV1*, double*, int32_t);
+extern "C" int32_t qwen36_pipeline_abort_v1(void*);
+extern "C" void qwen36_set_checkpoint(void*, int32_t, int64_t);
+extern "C" int32_t qwen36_set_max_grad_norm(void*, double);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int64_t qwen36_add_lora_v2(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" void* qwen36_get_adapter_optimizer_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t);
+extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t);
+extern "C" int64_t qwen36_get_lora_count(void*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" void* qwen36_get_lora_grad_accumulator(
+    void*, int64_t, int32_t);
+extern "C" int32_t qwen36_set_lora_tensor(
+    void*, int64_t, int32_t, void*);
+extern "C" int32_t qwen36_abort_gradient_accumulation(void*);
+extern "C" int64_t qwen36_export_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" int64_t qwen36_get_step_count(void*);
+extern "C" int64_t qwen36_get_lora_batch_projection_build_count(void*);
+extern "C" void qwen36_free_training_context(void*);
+
+static at::Tensor cuda_tensor(std::initializer_list<int64_t> shape, double value) {
+    auto options = at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16);
+    return at::full(shape, value, options);
+}
+
+static at::Tensor cuda_pattern(
+    std::initializer_list<int64_t> shape,
+    double scale,
+    double offset
+) {
+    int64_t numel = 1;
+    for (const auto dimension : shape) numel *= dimension;
+    auto options = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat);
+    return at::linspace(-scale, scale, numel, options)
+        .add(offset).reshape(shape).to(at::kBFloat16);
+}
+
+static at::Tensor& tensor_from_ptr(void* pointer) {
+    assert(pointer);
+    return *reinterpret_cast<at::Tensor*>(pointer);
+}
+
+static double max_diff(const at::Tensor& left, const at::Tensor& right) {
+    assert(left.sizes() == right.sizes());
+    return left.to(at::kFloat).sub(right.to(at::kFloat))
+        .abs().max().item<double>();
+}
+
+static double max_abs(const at::Tensor& tensor) {
+    return tensor.to(at::kFloat).abs().max().item<double>();
+}
+
+struct FixedState {
+    std::vector<at::Tensor> lora_a;
+    std::vector<at::Tensor> lora_b;
+    std::vector<at::Tensor> adam_m;
+    std::vector<at::Tensor> adam_v;
+    int64_t step = 0;
+};
+
+static FixedState snapshot_fixed_state(void* context) {
+    FixedState state;
+    const int64_t lora_count = qwen36_get_lora_count(context);
+    assert(lora_count > 0);
+    state.lora_a.reserve(lora_count);
+    state.lora_b.reserve(lora_count);
+    for (int64_t index = 0; index < lora_count; ++index) {
+        state.lora_a.push_back(tensor_from_ptr(
+            qwen36_get_lora_a(context, index)).detach().cpu().clone());
+        state.lora_b.push_back(tensor_from_ptr(
+            qwen36_get_lora_b(context, index)).detach().cpu().clone());
+    }
+
+    const int64_t optimizer_count = 2 * lora_count;
+    std::vector<void*> m_ptrs(optimizer_count);
+    std::vector<void*> v_ptrs(optimizer_count);
+    assert(qwen36_export_optimizer_state(
+        context, m_ptrs.data(), v_ptrs.data(), optimizer_count) ==
+        optimizer_count);
+    state.adam_m.reserve(optimizer_count);
+    state.adam_v.reserve(optimizer_count);
+    for (int64_t index = 0; index < optimizer_count; ++index) {
+        state.adam_m.push_back(
+            tensor_from_ptr(m_ptrs[index]).detach().cpu().clone());
+        state.adam_v.push_back(
+            tensor_from_ptr(v_ptrs[index]).detach().cpu().clone());
+    }
+    state.step = qwen36_get_step_count(context);
+    return state;
+}
+
+static void copy_fixed_lora(void* source, void* destination) {
+    const int64_t count = qwen36_get_lora_count(source);
+    assert(count == qwen36_get_lora_count(destination));
+    for (int64_t index = 0; index < count; ++index) {
+        assert(qwen36_set_lora_tensor(
+            destination, index, 0, qwen36_get_lora_a(source, index)) == 0);
+        assert(qwen36_set_lora_tensor(
+            destination, index, 1, qwen36_get_lora_b(source, index)) == 0);
+    }
+}
+
+struct StateDiff {
+    double parameter = 0.0;
+    double adam_m = 0.0;
+    double adam_v = 0.0;
+};
+
+static StateDiff compare_fixed_state(
+    const FixedState& left,
+    const FixedState& right
+) {
+    assert(left.lora_a.size() == right.lora_a.size());
+    assert(left.lora_b.size() == right.lora_b.size());
+    assert(left.adam_m.size() == right.adam_m.size());
+    assert(left.adam_v.size() == right.adam_v.size());
+    StateDiff difference;
+    for (size_t index = 0; index < left.lora_a.size(); ++index) {
+        difference.parameter = std::max(
+            difference.parameter,
+            max_diff(left.lora_a[index], right.lora_a[index]));
+        difference.parameter = std::max(
+            difference.parameter,
+            max_diff(left.lora_b[index], right.lora_b[index]));
+    }
+    for (size_t index = 0; index < left.adam_m.size(); ++index) {
+        difference.adam_m = std::max(
+            difference.adam_m,
+            max_diff(left.adam_m[index], right.adam_m[index]));
+        difference.adam_v = std::max(
+            difference.adam_v,
+            max_diff(left.adam_v[index], right.adam_v[index]));
+    }
+    return difference;
+}
+
+int main() {
+    assert(qwen36_kernel_abi_version() == 31);
+    const int rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0");
+    const int local_rank = std::atoi(
+        std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0");
+    const int pp_size = std::atoi(
+        std::getenv("PP_SIZE") ? std::getenv("PP_SIZE") : "2");
+    const int world_size = std::atoi(
+        std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1");
+    assert(pp_size >= 2 && world_size == pp_size);
+    assert(rank >= 0 && rank < pp_size);
+    c10::cuda::CUDAGuard guard(local_rank);
+
+    const int pp_rank = rank;
+    std::vector<at::Tensor> weights = {
+        cuda_tensor({4}, 1.0),
+        cuda_tensor({4}, 1.0),
+        cuda_pattern({8, 4}, 0.08, 0.02),
+        cuda_tensor({4}, 1.0),
+        cuda_pattern({4, 4}, 0.07, -0.01),
+        cuda_tensor({4}, 1.0),
+        cuda_pattern({4, 4}, 0.06, 0.015),
+        cuda_pattern({4, 4}, 0.08, 0.01),
+        cuda_pattern({8, 4}, 0.04, 0.02),
+        cuda_pattern({8, 4}, 0.05, -0.01),
+        cuda_pattern({4, 8}, 0.04, 0.01),
+    };
+    for (auto& weight : weights) weight.set_requires_grad(false);
+    auto embed = cuda_pattern({4, 4}, 0.20, 0.05);
+    auto final_norm = cuda_tensor({4}, 1.0);
+    auto lm_head = cuda_pattern({4, 4}, 0.15, -0.025);
+    embed.set_requires_grad(false);
+    final_norm.set_requires_grad(false);
+    lm_head.set_requires_grad(false);
+
+    std::vector<void*> weight_ptrs;
+    for (auto& weight : weights) weight_ptrs.push_back(&weight);
+    LayerConfig config{};
+    config.layer_type = 0;
+    config.num_heads = 1;
+    config.num_kv_heads = 1;
+    config.head_dim = 4;
+    config.partial_rotary_factor = 1.0;
+    config.rope_theta = 10000.0;
+    config.rms_eps = 1e-6;
+    config.intermediate_size = 8;
+    const int32_t stage_flags = pp_rank == 0 ? 1
+        : (pp_rank + 1 == pp_size ? 2 : 0);
+    // Keep the last PP stage free of active LoRA parameters. Its clipping
+    // path must still participate in the scalar PP norm reduction.
+    std::vector<int64_t> target_layers{0};
+    auto create_context = [&](bool synchronize_parameters) {
+        void* context = qwen36_create_training_context_v2(
+            weight_ptrs.data(), static_cast<int64_t>(weight_ptrs.size()),
+            pp_rank == 0 ? &embed : nullptr,
+            pp_rank + 1 == pp_size ? &final_norm : nullptr,
+            pp_rank + 1 == pp_size ? &lm_head : nullptr,
+            &config, 1, pp_rank, pp_size, stage_flags,
+            static_cast<int32_t>(at::kBFloat16),
+            1.0, 1e-3, 0.9, 0.999, 1e-8, 4, 1e-6, 1,
+            target_layers.data(), static_cast<int64_t>(target_layers.size()),
+            "q_proj", 0);
+        assert(context);
+        const auto init = synchronize_parameters
+            ? qwen36_init_parallel_nccl_v2(
+                context, rank, world_size,
+                0, 1, 0,
+                0, 1, 0,
+                0, 1, 0,
+                0, 1, 0,
+                pp_rank, pp_size, 0)
+            : qwen36_attach_parallel_nccl_no_sync_v2(
+                context, rank, world_size,
+                0, 1, 0,
+                0, 1, 0,
+                0, 1, 0,
+                0, 1, 0,
+                pp_rank, pp_size, 0);
+        assert(init == 0);
+        // Exercise the native null-mask path while keeping this fixture's
+        // token stream unchanged (99 is absent from the vocabulary rows).
+        assert(qwen36_set_pad_token_id(context, 99) == 0);
+        return context;
+    };
+    void* window_context = create_context(true);
+    void* legacy_context = create_context(false);
+    copy_fixed_lora(window_context, legacy_context);
+    const auto initial_state = snapshot_fixed_state(window_context);
+    const auto copied_state = snapshot_fixed_state(legacy_context);
+    const auto initial_difference = compare_fixed_state(
+        initial_state, copied_state);
+    assert(initial_difference.parameter == 0.0);
+    assert(initial_difference.adam_m == 0.0);
+    assert(initial_difference.adam_v == 0.0);
+
+    const int64_t sequence_length = std::atoll(
+        std::getenv("QWEN36_PP_SMOKE_SEQ") ?
+            std::getenv("QWEN36_PP_SMOKE_SEQ") : "3");
+    assert(sequence_length > 1);
+    auto long_options = at::TensorOptions().device(at::kCUDA).dtype(at::kLong);
+    auto token_positions = at::arange(sequence_length, long_options);
+    std::vector<at::Tensor> input_microbatches;
+    std::vector<at::Tensor> target_microbatches;
+    input_microbatches.reserve(4);
+    target_microbatches.reserve(4);
+    for (int64_t microbatch = 0; microbatch < 4; ++microbatch) {
+        input_microbatches.push_back(
+            token_positions.add(microbatch).remainder(4)
+                .reshape({1, sequence_length}));
+        auto target = at::ones({1, sequence_length}, long_options);
+        const int64_t masked_tokens = std::min(
+            microbatch, std::max<int64_t>(sequence_length - 2, 0));
+        if (masked_tokens > 0) {
+            target.narrow(
+                1, sequence_length - masked_tokens, masked_tokens).zero_();
+        }
+        target_microbatches.push_back(std::move(target));
+    }
+    const int64_t num_microbatches = 4;
+    const int64_t warmup = std::min<int64_t>(
+        pp_size - pp_rank - 1, num_microbatches);
+    std::vector<std::pair<int64_t, int64_t>> schedule;
+    for (int64_t microbatch = 0; microbatch < warmup; ++microbatch)
+        schedule.emplace_back(microbatch, -1);
+    for (int64_t microbatch = warmup;
+         microbatch < num_microbatches; ++microbatch)
+        schedule.emplace_back(microbatch, microbatch - warmup);
+    for (int64_t microbatch = num_microbatches - warmup;
+         microbatch < num_microbatches; ++microbatch)
+        schedule.emplace_back(-1, microbatch);
+
+    // Runtime kernel choices are part of the PP contract even when a stage
+    // does not own a GDN layer. Reject a mismatch before any data P2P begins.
+    if (rank == 0) {
+        setenv("QWEN36_GDN_INVERSE_BWD", "1", 1);
+    } else {
+        unsetenv("QWEN36_GDN_INVERSE_BWD");
+    }
+    Qwen36PipelineWindowV1 runtime_mismatch_window{
+        sizeof(Qwen36PipelineWindowV1), 1, 6, 1, 0, 1, 0};
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &runtime_mismatch_window) != 0);
+    unsetenv("QWEN36_GDN_INVERSE_BWD");
+
+    // Frozen FC1 fusion is a PP-wide runtime choice even when only one stage
+    // owns a dense or shared MLP.  Preserve the caller's setting after the
+    // negative test so the real window below uses the requested mode.
+    const bool fused_mlp_fc1_enabled =
+        std::getenv("QWEN36_FUSED_MLP_FC1") &&
+        std::atoi(std::getenv("QWEN36_FUSED_MLP_FC1")) != 0;
+    if (rank == 0) {
+        setenv("QWEN36_FUSED_MLP_FC1",
+            fused_mlp_fc1_enabled ? "0" : "1", 1);
+    } else {
+        setenv("QWEN36_FUSED_MLP_FC1",
+            fused_mlp_fc1_enabled ? "1" : "0", 1);
+    }
+    runtime_mismatch_window.window_id = 14;
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &runtime_mismatch_window) != 0);
+    if (fused_mlp_fc1_enabled) {
+        setenv("QWEN36_FUSED_MLP_FC1", "1", 1);
+    } else {
+        unsetenv("QWEN36_FUSED_MLP_FC1");
+    }
+
+    if (rank == 0) {
+        setenv("QWEN36_GDN_STATE_CHECKPOINT_STRIDE", "invalid", 1);
+    } else {
+        setenv("QWEN36_GDN_STATE_CHECKPOINT_STRIDE", "4", 1);
+    }
+    runtime_mismatch_window.window_id = 2;
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &runtime_mismatch_window) != 0);
+    unsetenv("QWEN36_GDN_STATE_CHECKPOINT_STRIDE");
+
+    if (rank == 0) qwen36_set_checkpoint(window_context, 1, 2);
+    runtime_mismatch_window.window_id = 3;
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &runtime_mismatch_window) != 0);
+    qwen36_set_checkpoint(window_context, 0, 1);
+
+    if (rank == 0) {
+        setenv("QWEN36_EP_A2A_PACKED", "0", 1);
+    } else {
+        setenv("QWEN36_EP_A2A_PACKED", "1", 1);
+    }
+    runtime_mismatch_window.window_id = 4;
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &runtime_mismatch_window) != 0);
+    unsetenv("QWEN36_EP_A2A_PACKED");
+
+    // The v1 implementation has one slot per microbatch and a canonical
+    // non-interleaved 1F1B schedule.  Reject unsupported scheduler/chunk
+    // encodings collectively instead of silently treating them as chunk 0.
+    Qwen36PipelineWindowV1 unsupported_schedule_window{
+        sizeof(Qwen36PipelineWindowV1), 1, 15, 1, 1, 1, 0};
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &unsupported_schedule_window) != 0);
+    unsupported_schedule_window.window_id = 16;
+    unsupported_schedule_window.schedule = 0;
+    unsupported_schedule_window.num_chunks = 2;
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &unsupported_schedule_window) != 0);
+
+    setenv("QWEN36_GDN_STATE_CHECKPOINT_STRIDE", "invalid", 1);
+    runtime_mismatch_window.window_id = 5;
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &runtime_mismatch_window) != 0);
+    unsetenv("QWEN36_GDN_STATE_CHECKPOINT_STRIDE");
+
+    setenv("QWEN36_GDN_STATE_CHECKPOINT_STRIDE", "0", 1);
+    runtime_mismatch_window.window_id = 10;
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &runtime_mismatch_window) != 0);
+    unsetenv("QWEN36_GDN_STATE_CHECKPOINT_STRIDE");
+
+    if (rank == 0) {
+        setenv("QWEN36_DISABLE_GROUPED_MM", "1", 1);
+    } else {
+        unsetenv("QWEN36_DISABLE_GROUPED_MM");
+    }
+    runtime_mismatch_window.window_id = 8;
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &runtime_mismatch_window) != 0);
+    unsetenv("QWEN36_DISABLE_GROUPED_MM");
+
+    if (rank == 0) {
+        setenv("QWEN36_CE_TOKEN_TILE", "128", 1);
+    } else {
+        setenv("QWEN36_CE_TOKEN_TILE", "256", 1);
+    }
+    runtime_mismatch_window.window_id = 9;
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &runtime_mismatch_window) != 0);
+    unsetenv("QWEN36_CE_TOKEN_TILE");
+
+    assert(qwen36_set_max_grad_norm(
+        window_context, rank == 0 ? 0.0 : 1e-4) == 0);
+    runtime_mismatch_window.window_id = 11;
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &runtime_mismatch_window) != 0);
+    assert(qwen36_set_max_grad_norm(window_context, 1e-4) == 0);
+    assert(qwen36_set_max_grad_norm(legacy_context, 1e-4) == 0);
+
+    Qwen36PipelineWindowV1 window{
+        sizeof(Qwen36PipelineWindowV1), 1, 7, num_microbatches, 0, 1, 0};
+    const int64_t window_builds_before =
+        qwen36_get_lora_batch_projection_build_count(window_context);
+    assert(qwen36_pipeline_begin_v1(window_context, &window) == 0);
+    // Registry identity must remain immutable for the whole PP window.  A
+    // failed mutation must be uniform across stages and must not abort the
+    // active window or desynchronize its later P2P schedule.
+    const int64_t dynamic_target_layer = 0;
+    assert(qwen36_add_lora(
+        window_context, 2, 4.0, &dynamic_target_layer, 1, "q_proj") < 0);
+    int64_t max_in_flight = 0;
+    Qwen36PipelineResultV1 result{};
+    result.struct_size = sizeof(Qwen36PipelineResultV1);
+    result.version = 1;
+    for (const auto& [forward_mb, backward_mb] : schedule) {
+        Qwen36PipelineTickV1 tick{
+            sizeof(Qwen36PipelineTickV1), 1, 7, forward_mb, backward_mb,
+            0, forward_mb < 0 ? 2 : (backward_mb < 0 ? 0 : 1),
+            forward_mb >= 0 ? &input_microbatches[forward_mb] : nullptr,
+            forward_mb >= 0 ? &target_microbatches[forward_mb] : nullptr,
+            nullptr, 0.25};
+        assert(qwen36_pipeline_tick_v1(window_context, &tick, &result) == 0);
+        max_in_flight = std::max(max_in_flight, result.in_flight);
+    }
+    assert(max_in_flight <= warmup);
+    assert(qwen36_get_step_count(window_context) == 0);
+    const int64_t window_projection_builds =
+        qwen36_get_lora_batch_projection_build_count(window_context) -
+        window_builds_before;
+    if (pp_rank == 0) assert(window_projection_builds > 0);
+    else assert(window_projection_builds == 0);
+
+    double gradient_difference = 0.0;
+    double gradient_magnitude = 0.0;
+    double legacy_probe_loss = 0.0;
+    const int64_t legacy_builds_before =
+        qwen36_get_lora_batch_projection_build_count(legacy_context);
+    for (int64_t microbatch = 0; microbatch < 4; ++microbatch) {
+        legacy_probe_loss = qwen36_train_micro_step(
+            legacy_context, &input_microbatches[microbatch],
+            &target_microbatches[microbatch], nullptr, 0.25, 0);
+        assert(std::isfinite(legacy_probe_loss) && legacy_probe_loss >= 0.0);
+    }
+    const int64_t legacy_projection_builds =
+        qwen36_get_lora_batch_projection_build_count(legacy_context) -
+        legacy_builds_before;
+    assert(legacy_projection_builds ==
+        window_projection_builds * num_microbatches);
+    const int64_t lora_count = qwen36_get_lora_count(window_context);
+    assert(lora_count == qwen36_get_lora_count(legacy_context));
+    for (int64_t index = 0; index < lora_count; ++index) {
+        for (int32_t is_b = 0; is_b <= 1; ++is_b) {
+            auto* window_gradient = reinterpret_cast<at::Tensor*>(
+                qwen36_get_lora_grad_accumulator(
+                    window_context, index, is_b));
+            auto* legacy_gradient = reinterpret_cast<at::Tensor*>(
+                qwen36_get_lora_grad_accumulator(
+                    legacy_context, index, is_b));
+            assert((window_gradient == nullptr) == (legacy_gradient == nullptr));
+            if (!window_gradient) continue;
+            gradient_difference = std::max(
+                gradient_difference,
+                max_diff(*window_gradient, *legacy_gradient));
+            gradient_magnitude = std::max(
+                gradient_magnitude, max_abs(*window_gradient));
+        }
+    }
+    if (pp_rank == 0) assert(gradient_magnitude > 1e-8);
+    else assert(gradient_magnitude == 0.0);
+    assert(gradient_difference <= 1e-6);
+    assert(qwen36_abort_gradient_accumulation(legacy_context) == 0);
+    assert(qwen36_get_step_count(legacy_context) == 0);
+
+    assert(qwen36_pipeline_finish_v1(window_context, 1, &result) == 0);
+    assert(result.completed_fwd == 4 && result.completed_bwd == 4);
+    assert(result.in_flight == 0 && std::isfinite(result.loss));
+    assert(qwen36_get_step_count(window_context) == 1);
+    const double window_loss = result.loss;
+
+    double legacy_loss_sum = 0.0;
+    for (int64_t microbatch = 0; microbatch < 4; ++microbatch) {
+        const double legacy_microbatch_loss = qwen36_train_micro_step(
+            legacy_context, &input_microbatches[microbatch],
+            &target_microbatches[microbatch], nullptr, 0.25,
+            microbatch == 3 ? 1 : 0);
+        assert(std::isfinite(legacy_microbatch_loss) &&
+            legacy_microbatch_loss >= 0.0);
+        legacy_loss_sum += legacy_microbatch_loss;
+    }
+    assert(qwen36_get_step_count(legacy_context) == 1);
+    const double legacy_loss = legacy_loss_sum / 4.0;
+    assert(std::abs(window_loss - legacy_loss) <= 1e-5);
+
+    const auto window_state = snapshot_fixed_state(window_context);
+    const auto legacy_state = snapshot_fixed_state(legacy_context);
+    const auto parity = compare_fixed_state(window_state, legacy_state);
+    assert(window_state.step == 1 && legacy_state.step == 1);
+    assert(parity.parameter <= 2e-3);
+    assert(parity.adam_m <= 1e-5);
+    assert(parity.adam_v <= 1e-7);
+    const auto update = compare_fixed_state(initial_state, window_state);
+    if (pp_rank == 0) {
+        assert(update.parameter > 0.0);
+        assert(std::max(update.adam_m, update.adam_v) > 0.0);
+    } else {
+        assert(update.parameter == 0.0);
+        assert(update.adam_m == 0.0 && update.adam_v == 0.0);
+    }
+    assert(update.adam_m <= 1.1e-5);
+
+    auto bad_target_mask = at::ones(
+        {1, sequence_length - 1},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong));
+    Qwen36PipelineWindowV1 late_bad_window{
+        sizeof(Qwen36PipelineWindowV1), 1, 8, 2, 0, 1, 0};
+    assert(qwen36_pipeline_begin_v1(window_context, &late_bad_window) == 0);
+    const int64_t late_warmup = std::min<int64_t>(pp_size - pp_rank - 1, 2);
+    std::vector<std::pair<int64_t, int64_t>> late_schedule;
+    for (int64_t microbatch = 0; microbatch < late_warmup; ++microbatch)
+        late_schedule.emplace_back(microbatch, -1);
+    for (int64_t microbatch = late_warmup; microbatch < 2; ++microbatch)
+        late_schedule.emplace_back(microbatch, microbatch - late_warmup);
+    for (int64_t microbatch = 2 - late_warmup; microbatch < 2; ++microbatch)
+        late_schedule.emplace_back(-1, microbatch);
+    for (const auto& [forward_mb, backward_mb] : late_schedule) {
+        void* target = nullptr;
+        if (forward_mb >= 0) {
+            target = forward_mb == 1 && rank == 0
+                ? static_cast<void*>(&bad_target_mask)
+                : static_cast<void*>(&target_microbatches[forward_mb]);
+        }
+        int32_t phase = forward_mb < 0 ? 2 : (backward_mb < 0 ? 0 : 1);
+        if (forward_mb == 1 && rank == 1) phase = (phase + 1) % 3;
+        Qwen36PipelineTickV1 tick{
+            sizeof(Qwen36PipelineTickV1), 1, 8, forward_mb, backward_mb,
+            0, phase,
+            forward_mb >= 0 ? &input_microbatches[forward_mb] : nullptr,
+            target, nullptr, 0.5};
+        assert(qwen36_pipeline_tick_v1(window_context, &tick, &result) == 0);
+    }
+    assert(qwen36_pipeline_finish_v1(window_context, 1, &result) != 0);
+    assert(qwen36_get_step_count(window_context) == 1);
+
+    // A first-tick metadata error must still participate in the initial
+    // contract handshake.  The window should drain its fixed P2P schedule and
+    // fail uniformly at finish instead of deadlocking on the control stream.
+    Qwen36PipelineWindowV1 first_metadata_bad_window{
+        sizeof(Qwen36PipelineWindowV1), 1, 9, 2, 0, 1, 0};
+    assert(qwen36_pipeline_begin_v1(
+        window_context, &first_metadata_bad_window) == 0);
+    const int64_t first_warmup = std::min<int64_t>(
+        pp_size - pp_rank - 1, 2);
+    std::vector<std::pair<int64_t, int64_t>> first_metadata_schedule;
+    for (int64_t microbatch = 0; microbatch < first_warmup; ++microbatch)
+        first_metadata_schedule.emplace_back(microbatch, -1);
+    for (int64_t microbatch = first_warmup; microbatch < 2; ++microbatch)
+        first_metadata_schedule.emplace_back(
+            microbatch, microbatch - first_warmup);
+    for (int64_t microbatch = 2 - first_warmup; microbatch < 2; ++microbatch)
+        first_metadata_schedule.emplace_back(-1, microbatch);
+    for (const auto& [forward_mb, backward_mb] : first_metadata_schedule) {
+        int32_t phase = forward_mb < 0 ? 2 : (backward_mb < 0 ? 0 : 1);
+        if (forward_mb == 0 && rank == 1) phase = (phase + 1) % 3;
+        Qwen36PipelineTickV1 first_metadata_bad_tick{
+            sizeof(Qwen36PipelineTickV1), 1, 9, forward_mb, backward_mb,
+            0, phase,
+            forward_mb >= 0 ? &input_microbatches[forward_mb] : nullptr,
+            forward_mb >= 0 ? &target_microbatches[forward_mb] : nullptr,
+            nullptr, 1.0};
+        assert(qwen36_pipeline_tick_v1(
+            window_context, &first_metadata_bad_tick, &result) == 0);
+    }
+    assert(qwen36_pipeline_finish_v1(window_context, 1, &result) != 0);
+    assert(qwen36_get_step_count(window_context) == 1);
+
+    // Deliberately make only rank 0 violate the next window's shape contract.
+    // All ranks must fail at the PP min/max preflight, before any NCCL P2P
+    // count can diverge.
+    auto* bad_target_ptr = rank == 0
+        ? &bad_target_mask : &target_microbatches[0];
+    Qwen36PipelineWindowV1 bad_window{
+        sizeof(Qwen36PipelineWindowV1), 1, 10, 1, 0, 1, 0};
+    assert(qwen36_pipeline_begin_v1(window_context, &bad_window) == 0);
+    const int64_t bad_backward = warmup == 0 ? 0 : -1;
+    Qwen36PipelineTickV1 bad_tick{
+        sizeof(Qwen36PipelineTickV1), 1, 10, 0, bad_backward, 0,
+        bad_backward < 0 ? 0 : 1,
+        &input_microbatches[0], bad_target_ptr, nullptr, 1.0};
+    assert(qwen36_pipeline_tick_v1(window_context, &bad_tick, &result) != 0);
+
+    const int64_t adapter_one = qwen36_add_lora(
+        window_context, 2, 4.0, &dynamic_target_layer, 1, "q_proj");
+    const int64_t adapter_two = qwen36_add_lora_v2(
+        window_context, 3, 4.0, &dynamic_target_layer, 1, "q_proj");
+    const int64_t adapter_unselected = qwen36_add_lora(
+        window_context, 2, 4.0, &dynamic_target_layer, 1, "q_proj");
+    assert(adapter_one > 0 && adapter_two > adapter_one &&
+        adapter_unselected > adapter_two);
+    assert(qwen36_get_adapter_step_count(window_context, adapter_one) == 0);
+    assert(qwen36_get_adapter_step_count(window_context, adapter_two) == 0);
+    assert(qwen36_get_adapter_step_count(
+        window_context, adapter_unselected) == 0);
+
+    Qwen36PipelineWindowV1 mismatched_dynamic_window{
+        sizeof(Qwen36PipelineWindowV1), 1, 12, 2, 0, 1, 1};
+    const int64_t mismatched_ids[] = {
+        rank == 0 ? adapter_one : adapter_two,
+        rank == 0 ? adapter_two : adapter_one,
+    };
+    assert(qwen36_pipeline_begin_dynamic_selected_v1(
+        window_context, &mismatched_dynamic_window,
+        mismatched_ids, 2) != 0);
+
+    auto* selected_one_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            window_context, adapter_one, 0, "q_proj", 1));
+    auto* selected_two_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            window_context, adapter_two, 0, "q_proj", 1));
+    auto* unselected_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            window_context, adapter_unselected, 0, "q_proj", 1));
+    auto* unselected_m = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_optimizer_tensor(
+            window_context, adapter_unselected, 0, "q_proj", 1, 0));
+    auto* unselected_v = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_optimizer_tensor(
+            window_context, adapter_unselected, 0, "q_proj", 1, 1));
+    assert((selected_one_b != nullptr) == (pp_rank == 0));
+    assert((selected_two_b != nullptr) == (pp_rank == 0));
+    assert((unselected_b != nullptr) == (pp_rank == 0));
+    assert((unselected_m != nullptr) == (pp_rank == 0));
+    assert((unselected_v != nullptr) == (pp_rank == 0));
+    at::Tensor selected_one_before;
+    at::Tensor selected_two_before;
+    at::Tensor unselected_before;
+    at::Tensor unselected_m_before;
+    at::Tensor unselected_v_before;
+    if (pp_rank == 0) {
+        selected_one_before = selected_one_b->clone();
+        selected_two_before = selected_two_b->clone();
+        unselected_before = unselected_b->clone();
+        unselected_m_before = unselected_m->clone();
+        unselected_v_before = unselected_v->clone();
+    }
+
+    const int64_t dynamic_sequence_length =
+        std::max<int64_t>(sequence_length, 3);
+    auto dynamic_positions = at::arange(dynamic_sequence_length, long_options);
+    std::vector<at::Tensor> dynamic_inputs;
+    std::vector<at::Tensor> dynamic_targets;
+    for (int64_t microbatch = 0; microbatch < 2; ++microbatch) {
+        dynamic_inputs.push_back(at::stack({
+            dynamic_positions.add(microbatch).remainder(4),
+            dynamic_positions.add(microbatch + 2).remainder(4),
+        }));
+        auto target = at::ones({2, dynamic_sequence_length}, long_options);
+        target.select(0, 1).select(0, dynamic_sequence_length - 1).zero_();
+        dynamic_targets.push_back(std::move(target));
+    }
+    const int64_t dynamic_microbatches = 2;
+    const int64_t dynamic_warmup = std::min<int64_t>(
+        pp_size - pp_rank - 1, dynamic_microbatches);
+    std::vector<std::pair<int64_t, int64_t>> dynamic_schedule;
+    for (int64_t microbatch = 0; microbatch < dynamic_warmup; ++microbatch)
+        dynamic_schedule.emplace_back(microbatch, -1);
+    for (int64_t microbatch = dynamic_warmup;
+         microbatch < dynamic_microbatches; ++microbatch)
+        dynamic_schedule.emplace_back(microbatch, microbatch - dynamic_warmup);
+    for (int64_t microbatch = dynamic_microbatches - dynamic_warmup;
+         microbatch < dynamic_microbatches; ++microbatch)
+        dynamic_schedule.emplace_back(-1, microbatch);
+
+    Qwen36PipelineWindowV1 dynamic_window{
+        sizeof(Qwen36PipelineWindowV1), 1, 13,
+        dynamic_microbatches, 0, 1, 1};
+    const int64_t selected_ids[] = {adapter_one, adapter_two};
+    assert(qwen36_pipeline_begin_dynamic_selected_v1(
+        window_context, &dynamic_window, selected_ids, 2) == 0);
+    for (const auto& [forward_mb, backward_mb] : dynamic_schedule) {
+        Qwen36PipelineTickV1 tick{
+            sizeof(Qwen36PipelineTickV1), 1, 13, forward_mb, backward_mb,
+            0, forward_mb < 0 ? 2 : (backward_mb < 0 ? 0 : 1),
+            forward_mb >= 0 ? &dynamic_inputs[forward_mb] : nullptr,
+            forward_mb >= 0 ? &dynamic_targets[forward_mb] : nullptr,
+            nullptr, 0.5};
+        assert(qwen36_pipeline_tick_v1(
+            window_context, &tick, &result) == 0);
+    }
+    double tenant_losses[2] = {-1.0, -1.0};
+    assert(qwen36_pipeline_finish_dynamic_report_v1(
+        window_context, 1, &result, tenant_losses, 2) == 0);
+    assert(result.completed_fwd == dynamic_microbatches &&
+        result.completed_bwd == dynamic_microbatches &&
+        result.in_flight == 0 && result.optimizer_step == -1);
+    assert(std::isfinite(result.loss) && result.loss >= 0.0);
+    assert(std::isfinite(tenant_losses[0]) && tenant_losses[0] >= 0.0);
+    assert(std::isfinite(tenant_losses[1]) && tenant_losses[1] >= 0.0);
+    const double tenant_one_tokens =
+        static_cast<double>(dynamic_sequence_length - 1);
+    const double tenant_two_tokens =
+        static_cast<double>(dynamic_sequence_length - 2);
+    const double reported_aggregate =
+        (tenant_losses[0] * tenant_one_tokens +
+         tenant_losses[1] * tenant_two_tokens) /
+        (tenant_one_tokens + tenant_two_tokens);
+    assert(std::abs(result.loss - reported_aggregate) <= 1e-5);
+    assert(qwen36_get_adapter_step_count(window_context, adapter_one) == 1);
+    assert(qwen36_get_adapter_step_count(window_context, adapter_two) == 1);
+    assert(qwen36_get_adapter_step_count(
+        window_context, adapter_unselected) == 0);
+    if (pp_rank == 0) {
+        assert(max_diff(*selected_one_b, selected_one_before) > 0.0);
+        assert(max_diff(*selected_two_b, selected_two_before) > 0.0);
+        assert(max_diff(*unselected_b, unselected_before) == 0.0);
+        assert(max_diff(*unselected_m, unselected_m_before) == 0.0);
+        assert(max_diff(*unselected_v, unselected_v_before) == 0.0);
+    }
+
+    // A late dynamic metadata error must keep the selected-tenant batch
+    // contract (two rows here), even though the registry also contains an
+    // unselected third tenant.  This drains the fixed P2P schedule and fails
+    // uniformly at finish without a cross-stage count mismatch.
+    auto dynamic_bad_target = at::ones(
+        {1, dynamic_sequence_length - 1}, long_options);
+    Qwen36PipelineWindowV1 dynamic_late_bad_window{
+        sizeof(Qwen36PipelineWindowV1), 1, 14, 2, 0, 1, 1};
+    assert(qwen36_pipeline_begin_dynamic_selected_v1(
+        window_context, &dynamic_late_bad_window, selected_ids, 2) == 0);
+    const int64_t late_dynamic_warmup = std::min<int64_t>(
+        pp_size - pp_rank - 1, 2);
+    std::vector<std::pair<int64_t, int64_t>> late_dynamic_schedule;
+    for (int64_t microbatch = 0; microbatch < late_dynamic_warmup; ++microbatch)
+        late_dynamic_schedule.emplace_back(microbatch, -1);
+    for (int64_t microbatch = late_dynamic_warmup;
+         microbatch < 2; ++microbatch)
+        late_dynamic_schedule.emplace_back(microbatch,
+            microbatch - late_dynamic_warmup);
+    for (int64_t microbatch = 2 - late_dynamic_warmup;
+         microbatch < 2; ++microbatch)
+        late_dynamic_schedule.emplace_back(-1, microbatch);
+    for (const auto& [forward_mb, backward_mb] : late_dynamic_schedule) {
+        void* target = nullptr;
+        if (forward_mb >= 0) {
+            target = (forward_mb == 1 && rank == 0)
+                ? static_cast<void*>(&dynamic_bad_target)
+                : static_cast<void*>(&dynamic_targets[forward_mb]);
+        }
+        Qwen36PipelineTickV1 tick{
+            sizeof(Qwen36PipelineTickV1), 1, 14, forward_mb, backward_mb,
+            0, forward_mb < 0 ? 2 : (backward_mb < 0 ? 0 : 1),
+            forward_mb >= 0 ? &dynamic_inputs[forward_mb] : nullptr,
+            target, nullptr, 1.0};
+        assert(qwen36_pipeline_tick_v1(
+            window_context, &tick, &result) == 0);
+    }
+    assert(qwen36_pipeline_finish_dynamic_report_v1(
+        window_context, 1, &result, tenant_losses, 2) != 0);
+    assert(qwen36_get_adapter_step_count(window_context, adapter_one) == 1);
+    assert(qwen36_get_adapter_step_count(window_context, adapter_two) == 1);
+    assert(qwen36_get_adapter_step_count(
+        window_context, adapter_unselected) == 0);
+    std::printf("native_qwen36_pp_train rank=%d loss=%0.6f "
+        "grad_diff=%0.8e param_diff=%0.8e m_diff=%0.8e v_diff=%0.8e "
+        "max_in_flight=%ld cache_builds=%ld legacy_cache_builds=%ld "
+        "dynamic_loss=%0.6f tenant_losses=%0.6f,%0.6f step=1 ok\n", rank,
+        window_loss, gradient_difference, parity.parameter, parity.adam_m,
+        parity.adam_v, (long)max_in_flight, (long)window_projection_builds,
+        (long)legacy_projection_builds, result.loss,
+        tenant_losses[0], tenant_losses[1]);
+    qwen36_free_training_context(legacy_context);
+    qwen36_free_training_context(window_context);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_sequence_parallel_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_sequence_parallel_smoke.cpp
new file mode 100644
index 00000000..0a4a24e6
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_sequence_parallel_smoke.cpp
@@ -0,0 +1,404 @@
+#include <ATen/ATen.h>
+
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <initializer_list>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" int64_t qwen36_kernel_abi_version();
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_parallel_nccl(
+    void*, int32_t, int32_t, int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t, int32_t, int32_t, int32_t);
+extern "C" int32_t qwen36_set_pad_token_id(void*, int64_t);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" int64_t qwen36_get_sequence_parallel_counter(void*, int32_t);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" void* qwen36_get_adapter_optimizer_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t);
+extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t);
+extern "C" int32_t qwen36_train_multi_lora_selected_v3(
+    void*, void*, void*, void*, const int64_t*, int32_t,
+    double*, double*, int32_t);
+extern "C" void qwen36_free_training_context(void*);
+
+static at::Tensor values(std::initializer_list<int64_t> shape, double scale) {
+    int64_t count = 1;
+    for (auto dim : shape) count *= dim;
+    return ((at::arange(count, at::TensorOptions().device(at::kCUDA)
+        .dtype(at::kFloat)).remainder(23) - 11.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static std::vector<void*> ptrs(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static double max_diff(const at::Tensor& left, const at::Tensor& right) {
+    assert(left.sizes() == right.sizes());
+    return left.to(at::kFloat).sub(right.to(at::kFloat)).abs().max().item<double>();
+}
+
+static void set_dynamic_pair(
+    void* context, int64_t adapter, const char* module,
+    at::Tensor& a, at::Tensor& b
+) {
+    assert(qwen36_set_adapter_lora_tensor(
+        context, adapter, 0, module, 0, &a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        context, adapter, 0, module, 1, &b) == 0);
+}
+
+static void install_dynamic_fixture(
+    void* context, int64_t adapter, int rank, bool sharded, double scale
+) {
+    constexpr int64_t hidden = 8;
+    constexpr int64_t q_out = 16;
+    constexpr int64_t kv_out = 4;
+    constexpr int64_t intermediate = 12;
+    constexpr int64_t local_q_out = q_out / 2;
+    constexpr int64_t local_kv_out = kv_out / 2;
+    constexpr int64_t local_intermediate = intermediate / 2;
+    constexpr int64_t lora_rank = 2;
+
+    auto q_a = values({lora_rank, hidden}, scale);
+    auto q_b_full = values({q_out, lora_rank}, scale * 1.1);
+    auto k_a = values({lora_rank, hidden}, scale * 1.2);
+    auto k_b_full = values({kv_out, lora_rank}, scale * 1.3);
+    auto v_a = values({lora_rank, hidden}, scale * 1.4);
+    auto v_b_full = values({kv_out, lora_rank}, scale * 1.5);
+    auto o_a_full = values({lora_rank, hidden}, scale * 1.6);
+    auto o_b = values({hidden, lora_rank}, scale * 1.7);
+    auto gate_a = values({lora_rank, hidden}, scale * 1.8);
+    auto gate_b_full = values({intermediate, lora_rank}, scale * 1.9);
+    auto up_a = values({lora_rank, hidden}, scale * 2.0);
+    auto up_b_full = values({intermediate, lora_rank}, scale * 2.1);
+    auto down_a_full = values({lora_rank, intermediate}, scale * 2.2);
+    auto down_b = values({hidden, lora_rank}, scale * 2.3);
+
+    auto shard_rows = [rank](const at::Tensor& tensor, int64_t rows) {
+        return tensor.narrow(0, rank * rows, rows).contiguous();
+    };
+    auto shard_cols = [rank](const at::Tensor& tensor, int64_t cols) {
+        return tensor.narrow(1, rank * cols, cols).contiguous();
+    };
+    auto q_b = sharded ? shard_rows(q_b_full, local_q_out) : q_b_full;
+    auto k_b = sharded ? shard_rows(k_b_full, local_kv_out) : k_b_full;
+    auto v_b = sharded ? shard_rows(v_b_full, local_kv_out) : v_b_full;
+    auto o_a = sharded ? shard_cols(o_a_full, hidden / 2) : o_a_full;
+    auto gate_b = sharded
+        ? shard_rows(gate_b_full, local_intermediate) : gate_b_full;
+    auto up_b = sharded
+        ? shard_rows(up_b_full, local_intermediate) : up_b_full;
+    auto down_a = sharded
+        ? shard_cols(down_a_full, local_intermediate) : down_a_full;
+
+    set_dynamic_pair(context, adapter, "q_proj", q_a, q_b);
+    set_dynamic_pair(context, adapter, "k_proj", k_a, k_b);
+    set_dynamic_pair(context, adapter, "v_proj", v_a, v_b);
+    set_dynamic_pair(context, adapter, "o_proj", o_a, o_b);
+    set_dynamic_pair(context, adapter, "gate_proj", gate_a, gate_b);
+    set_dynamic_pair(context, adapter, "up_proj", up_a, up_b);
+    set_dynamic_pair(context, adapter, "down_proj", down_a, down_b);
+}
+
+int main() {
+    assert(qwen36_kernel_abi_version() == 31);
+    const int rank = std::atoi(std::getenv("RANK"));
+    const int world = std::atoi(std::getenv("WORLD_SIZE"));
+    const int local_rank = std::atoi(std::getenv("LOCAL_RANK"));
+    assert(world == 2 && (rank == 0 || rank == 1));
+    qwen36_set_cuda_device(local_rank);
+    setenv("QWEN36_SEQUENCE_PARALLEL", "1", 1);
+    setenv("TP_SIZE", "2", 1);
+    setenv("CP_SIZE", "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("PP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", std::to_string(rank).c_str(), 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+
+    constexpr int64_t hidden = 8, heads = 4, kv_heads = 2, head_dim = 2;
+    constexpr int64_t intermediate = 12, vocab = 16, rank_lora = 2;
+    constexpr int64_t local_heads = heads / 2, local_kv = kv_heads / 2;
+    std::vector<at::Tensor> full;
+    full.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full.push_back(values({2 * heads * head_dim, hidden}, .01));
+    full.push_back(at::ones({head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full.push_back(values({kv_heads * head_dim, hidden}, .012));
+    full.push_back(at::ones({head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full.push_back(values({kv_heads * head_dim, hidden}, .008));
+    full.push_back(values({hidden, heads * head_dim}, .011));
+    full.push_back(values({intermediate, hidden}, .009));
+    full.push_back(values({intermediate, hidden}, .007));
+    full.push_back(values({hidden, intermediate}, .010));
+    std::vector<at::Tensor> local = {
+        full[0], full[1], full[2].narrow(0, rank * local_heads * 2 * head_dim,
+            local_heads * 2 * head_dim).contiguous(), full[3],
+        full[4].narrow(0, rank * local_kv * head_dim, local_kv * head_dim).contiguous(),
+        full[5], full[6].narrow(0, rank * local_kv * head_dim, local_kv * head_dim).contiguous(),
+        full[7].narrow(1, rank * local_heads * head_dim, local_heads * head_dim).contiguous(),
+        full[8].narrow(0, rank * intermediate / 2, intermediate / 2).contiguous(),
+        full[9].narrow(0, rank * intermediate / 2, intermediate / 2).contiguous(),
+        full[10].narrow(1, rank * intermediate / 2, intermediate / 2).contiguous(),
+    };
+    auto embed = values({vocab, hidden}, .02);
+    auto final_norm = at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+    auto lm_head = values({vocab, hidden}, .015);
+    auto local_embed = embed.narrow(0, rank * vocab / 2, vocab / 2).contiguous();
+    auto local_lm_head = lm_head.narrow(0, rank * vocab / 2, vocab / 2).contiguous();
+    LayerConfig config{};
+    config.num_heads = heads; config.num_kv_heads = kv_heads; config.head_dim = head_dim;
+    config.partial_rotary_factor = 1.0; config.rope_theta = 10000.0;
+    config.rms_eps = 1e-5; config.intermediate_size = intermediate;
+    const int64_t target_layer = 0;
+    auto local_ptrs = ptrs(local);
+    constexpr int32_t flags = (1 << 0) | (1 << 2) | (1 << 4);
+    void* ctx = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head, &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, vocab, 1e-5, rank_lora,
+        &target_layer, 1, "q_proj,k_proj,v_proj,o_proj", flags);
+    assert(ctx);
+    auto q_a = values({rank_lora, hidden}, .002);
+    auto q_b = values({local_heads * 2 * head_dim, rank_lora}, .001);
+    auto k_a = values({rank_lora, hidden}, .002);
+    auto k_b = values({local_kv * head_dim, rank_lora}, .001);
+    auto v_a = values({rank_lora, hidden}, .002);
+    auto v_b = values({local_kv * head_dim, rank_lora}, .001);
+    auto o_a = values({rank_lora, local_heads * head_dim}, .002);
+    auto o_b = values({hidden, rank_lora}, .001);
+    at::Tensor* factors[] = {&q_a, &q_b, &k_a, &k_b, &v_a, &v_b, &o_a, &o_b};
+    for (int64_t slot = 0; slot < 4; ++slot) {
+        assert(qwen36_set_lora_tensor(ctx, slot, 0, factors[2 * slot]) == 0);
+        assert(qwen36_set_lora_tensor(ctx, slot, 1, factors[2 * slot + 1]) == 0);
+    }
+    assert(qwen36_init_parallel_nccl(ctx, rank, world, rank, 2, 0,
+        0, 1, 0, 0, 1, 0) == 0);
+    assert(qwen36_set_pad_token_id(ctx, 0) == 0);
+    auto* b_before_ptr = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(ctx, 0));
+    assert(b_before_ptr);
+    auto b_before = b_before_ptr->clone();
+    auto ids = at::tensor({1, 2, 3, 4}, at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 4});
+    auto target = at::ones({1, 4}, at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    const double loss = qwen36_train_step(ctx, &ids, &target, nullptr);
+    assert(std::isfinite(loss) && loss > 0.0);
+    auto* b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(ctx, 0));
+    assert(b && (*b - b_before).abs().sum().item<double>() > 0.0);
+    assert(qwen36_get_sequence_parallel_counter(ctx, 0) == 1);
+    assert(qwen36_get_sequence_parallel_counter(ctx, 1) > 0);
+    assert(qwen36_get_sequence_parallel_counter(ctx, 2) > 0);
+    assert(qwen36_get_sequence_parallel_counter(ctx, 3) == 1);
+    std::printf("native_qwen36_sequence_parallel_smoke rank=%d loss=%0.8f ag=%lld rs=%lld local_seq=%lld\n",
+        rank, loss, (long long)qwen36_get_sequence_parallel_counter(ctx, 1),
+        (long long)qwen36_get_sequence_parallel_counter(ctx, 2),
+        (long long)qwen36_get_sequence_parallel_counter(ctx, 4));
+
+    // Dynamic multi-LoRA + sequence parallel oracle. Two selected tenants
+    // receive rows with different token counts; the third tenant must remain
+    // bitwise unchanged. The reference context owns the complete TP weights.
+    auto full_ptrs = ptrs(full);
+    const char* dynamic_targets =
+        "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj";
+    const int64_t dynamic_target_layer = 0;
+    void* dynamic_distributed = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head, &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, vocab, 1e-5, rank_lora,
+        &dynamic_target_layer, 1, dynamic_targets, flags);
+    assert(dynamic_distributed);
+    assert(qwen36_init_parallel_nccl(
+        dynamic_distributed, rank, world, rank, 2, 0,
+        0, 1, 0, 0, 1, 0) == 0);
+    assert(qwen36_set_pad_token_id(dynamic_distributed, 0) == 0);
+
+    unsetenv("QWEN36_SEQUENCE_PARALLEL");
+    setenv("WORLD_SIZE", "1", 1);
+    setenv("RANK", "0", 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    void* dynamic_reference = qwen36_create_training_context_ex(
+        full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, .9, .999, 1e-8, vocab, 1e-5, rank_lora,
+        &dynamic_target_layer, 1, dynamic_targets, 0);
+    assert(dynamic_reference);
+    assert(qwen36_set_pad_token_id(dynamic_reference, 0) == 0);
+    setenv("QWEN36_SEQUENCE_PARALLEL", "1", 1);
+    setenv("WORLD_SIZE", "2", 1);
+    setenv("RANK", std::to_string(rank).c_str(), 1);
+    setenv("TP_SIZE", "2", 1);
+    setenv("RUSTRAIN_TP_RANK", std::to_string(rank).c_str(), 1);
+
+    std::array<int64_t, 3> distributed_adapters{};
+    std::array<int64_t, 3> reference_adapters{};
+    for (int64_t index = 0; index < 3; ++index) {
+        distributed_adapters[index] = qwen36_add_lora(
+            dynamic_distributed, rank_lora, rank_lora,
+            &dynamic_target_layer, 1, dynamic_targets);
+        reference_adapters[index] = qwen36_add_lora(
+            dynamic_reference, rank_lora, rank_lora,
+            &dynamic_target_layer, 1, dynamic_targets);
+        assert(distributed_adapters[index] > 0 &&
+            reference_adapters[index] > 0);
+        install_dynamic_fixture(
+            dynamic_distributed, distributed_adapters[index], rank, true,
+            0.0008 + index * 0.00017);
+        install_dynamic_fixture(
+            dynamic_reference, reference_adapters[index], 0, false,
+            0.0008 + index * 0.00017);
+    }
+
+    auto dynamic_ids = at::tensor(
+        {1, 2, 3, 4, 5, 6, 7, 8},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({2, 4});
+    auto dynamic_targets_mask = at::tensor(
+        {1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)).reshape({2, 4});
+    const int64_t selected_dynamic_ids[] = {
+        distributed_adapters[0], distributed_adapters[1]};
+    const int64_t selected_reference_ids[] = {
+        reference_adapters[0], reference_adapters[1]};
+    double distributed_dynamic_loss = -1.0;
+    double reference_dynamic_loss = -1.0;
+    std::array<double, 2> distributed_tenant_losses{-1.0, -1.0};
+    std::array<double, 2> reference_tenant_losses{-1.0, -1.0};
+    const std::array<const char*, 7> dynamic_modules = {
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj"};
+    auto snapshot = [](void* context, int64_t adapter, const char* module) {
+        std::array<at::Tensor, 6> state;
+        for (int is_b = 0; is_b < 2; ++is_b) {
+            auto* parameter = reinterpret_cast<at::Tensor*>(
+                qwen36_get_adapter_lora_tensor(
+                    context, adapter, 0, module, is_b));
+            auto* m = reinterpret_cast<at::Tensor*>(
+                qwen36_get_adapter_optimizer_tensor(
+                    context, adapter, 0, module, is_b, 0));
+            auto* v = reinterpret_cast<at::Tensor*>(
+                qwen36_get_adapter_optimizer_tensor(
+                    context, adapter, 0, module, is_b, 1));
+            assert(parameter && m && v);
+            state[is_b] = parameter->clone();
+            state[2 + is_b] = m->clone();
+            state[4 + is_b] = v->clone();
+        }
+        return state;
+    };
+    std::array<std::array<at::Tensor, 6>, 7> unselected_before{};
+    for (size_t module = 0; module < dynamic_modules.size(); ++module)
+        unselected_before[module] = snapshot(
+            dynamic_distributed, distributed_adapters[2], dynamic_modules[module]);
+    assert(qwen36_train_multi_lora_selected_v3(
+        dynamic_distributed, &dynamic_ids, &dynamic_targets_mask,
+        nullptr, selected_dynamic_ids, 2,
+        &distributed_dynamic_loss, distributed_tenant_losses.data(), 2) == 0);
+    assert(qwen36_train_multi_lora_selected_v3(
+        dynamic_reference, &dynamic_ids, &dynamic_targets_mask,
+        nullptr, selected_reference_ids, 2,
+        &reference_dynamic_loss, reference_tenant_losses.data(), 2) == 0);
+    assert(std::isfinite(distributed_dynamic_loss) &&
+        std::isfinite(reference_dynamic_loss));
+
+    double dynamic_parameter_diff = 0.0;
+    double dynamic_m_diff = 0.0;
+    double dynamic_v_diff = 0.0;
+    for (int64_t tenant = 0; tenant < 2; ++tenant) {
+        for (size_t module_index = 0;
+             module_index < dynamic_modules.size(); ++module_index) {
+            const char* module = dynamic_modules[module_index];
+            const bool column_parallel = module_index != 3 && module_index != 6;
+            auto distributed = snapshot(
+                dynamic_distributed, distributed_adapters[tenant], module);
+            auto reference = snapshot(
+                dynamic_reference, reference_adapters[tenant], module);
+            const int64_t local_rows = module_index == 0 ? 8
+                : (module_index == 1 || module_index == 2 ? 2
+                : (module_index == 4 || module_index == 5 ? 6 : 8));
+            const int64_t local_cols = module_index == 3 ? 4
+                : (module_index == 6 ? 6 : 8);
+            for (int state = 0; state < 6; ++state) {
+                at::Tensor expected = reference[state];
+                if (column_parallel && (state == 1 || state == 3 || state == 5))
+                    expected = expected.narrow(0, rank * local_rows, local_rows);
+                if (!column_parallel && (state == 0 || state == 2 || state == 4))
+                    expected = expected.narrow(1, rank * local_cols, local_cols);
+                const double difference = max_diff(distributed[state], expected);
+                if (state == 0 || state == 1) dynamic_parameter_diff =
+                    std::max(dynamic_parameter_diff, difference);
+                else if (state == 2 || state == 3) dynamic_m_diff =
+                    std::max(dynamic_m_diff, difference);
+                else dynamic_v_diff = std::max(dynamic_v_diff, difference);
+            }
+        }
+    }
+    double unselected_diff = 0.0;
+    for (size_t module = 0; module < dynamic_modules.size(); ++module) {
+        auto after = snapshot(
+            dynamic_distributed, distributed_adapters[2], dynamic_modules[module]);
+        for (int state = 0; state < 6; ++state)
+            unselected_diff = std::max(
+                unselected_diff, max_diff(after[state], unselected_before[module][state]));
+    }
+    std::printf(
+        "native_qwen36_sequence_parallel_dynamic_raw rank=%d dist_loss=%0.8g ref_loss=%0.8g "
+        "dist_tenant=%0.8g,%0.8g ref_tenant=%0.8g,%0.8g\n",
+        rank, distributed_dynamic_loss, reference_dynamic_loss,
+        distributed_tenant_losses[0], distributed_tenant_losses[1],
+        reference_tenant_losses[0], reference_tenant_losses[1]);
+    std::fflush(stdout);
+    // TP/SP changes BF16 reduction order relative to the single-rank oracle.
+    // Keep the report tolerance wide enough for that rounding while state
+    // comparisons below remain substantially tighter.
+    assert(std::abs(distributed_dynamic_loss - reference_dynamic_loss) < 1e-2);
+    for (int tenant = 0; tenant < 2; ++tenant) {
+        assert(std::abs(distributed_tenant_losses[tenant] -
+            reference_tenant_losses[tenant]) < 1e-2);
+        assert(qwen36_get_adapter_step_count(
+            dynamic_distributed, distributed_adapters[tenant]) == 1);
+    }
+    assert(qwen36_get_adapter_step_count(
+        dynamic_distributed, distributed_adapters[2]) == 0);
+    assert(dynamic_parameter_diff < 5e-3 && dynamic_m_diff < 5e-4 &&
+        dynamic_v_diff < 5e-7 && unselected_diff == 0.0);
+    std::printf(
+        "native_qwen36_sequence_parallel_dynamic rank=%d loss_diff=%0.8g "
+        "tenant_loss_diff=%0.8g param_diff=%0.8g m_diff=%0.8g v_diff=%0.8g\n",
+        rank, std::abs(distributed_dynamic_loss - reference_dynamic_loss),
+        std::max(std::abs(distributed_tenant_losses[0] - reference_tenant_losses[0]),
+            std::abs(distributed_tenant_losses[1] - reference_tenant_losses[1])),
+        dynamic_parameter_diff, dynamic_m_diff, dynamic_v_diff);
+    qwen36_free_training_context(dynamic_reference);
+    qwen36_free_training_context(dynamic_distributed);
+    qwen36_free_training_context(ctx);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp
new file mode 100644
index 00000000..9f69bc8d
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp
@@ -0,0 +1,1398 @@
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include <cassert>
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include <sys/stat.h>
+#include <tuple>
+#include <unistd.h>
+#include <utility>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" int64_t qwen36_kernel_abi_version();
+extern "C" int32_t qwen36_init_nccl(void*);
+extern "C" void qwen36_set_nccl_comm(void*, void*, void*, int32_t, int32_t);
+extern "C" int64_t qwen36_get_lora_count(void*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t);
+extern "C" int32_t qwen36_abort_gradient_accumulation(void*);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" int32_t qwen36_set_router_aux_loss_coef(void*, double);
+extern "C" int32_t qwen36_set_max_grad_norm(void*, double);
+extern "C" void qwen36_set_checkpoint(void*, int32_t, int64_t);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_micro_step(
+    void*, void*, void*, void*, double, int32_t);
+extern "C" int64_t qwen36_get_step_count(void*);
+extern "C" int32_t qwen36_set_step_count(void*, int64_t);
+extern "C" int64_t qwen36_export_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t);
+extern "C" int32_t qwen36_validate_adapter_steps_v1(
+    void*, const int64_t*, const int64_t*, int32_t);
+extern "C" int64_t qwen36_get_dynamic_finalizer_count(void*);
+extern "C" int64_t qwen36_get_dynamic_adam_launch_count(void*);
+extern "C" int64_t qwen36_get_dynamic_train_batch_count(void*);
+extern "C" int32_t qwen36_get_accumulation_active(void*);
+extern "C" int32_t qwen36_get_context_health(void*);
+extern "C" double qwen36_get_accumulated_token_weight(void*);
+extern "C" int32_t qwen36_set_adapter_step_count(void*, int64_t, int64_t);
+extern "C" double qwen36_eval_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_multi_lora(
+    void*, void*, void*, void*, int32_t, int32_t);
+extern "C" double qwen36_train_multi_lora_selected(
+    void*, void*, void*, void*, const int64_t*, int32_t, int32_t);
+extern "C" double qwen36_train_multi_lora_selected_v2(
+    void*, void*, void*, void*, const int64_t*, int32_t);
+extern "C" double qwen36_train_step_host_i64(
+    void*, const int64_t*, const int64_t*, const int64_t*, int64_t, int64_t);
+extern "C" double qwen36_train_multi_lora_host_i64(
+    void*, const int64_t*, const int64_t*, const int64_t*, int64_t, int64_t,
+    int32_t, int32_t, const int64_t*, int32_t);
+extern "C" double qwen36_eval_step_host_i64(
+    void*, const int64_t*, const int64_t*, const int64_t*, int64_t, int64_t);
+extern "C" int32_t qwen36_eval_multi_lora_host_i64_v1(
+    void*, const int64_t*, const int64_t*, const int64_t*, int64_t, int64_t,
+    const int64_t*, int32_t, double*, int32_t);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int64_t qwen36_add_lora_v2(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int64_t qwen36_add_lora_with_optimizer(
+    void*, int64_t, double, const int64_t*, int64_t, const char*, double);
+extern "C" int64_t qwen36_add_lora_with_optimizer_v2(
+    void*, int64_t, double, const int64_t*, int64_t, const char*,
+    double, double, double, double);
+extern "C" int64_t qwen36_add_lora_for_restore(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int64_t qwen36_add_lora_for_restore_with_optimizer(
+    void*, int64_t, double, const int64_t*, int64_t, const char*, double);
+extern "C" int64_t qwen36_add_lora_for_restore_with_optimizer_v2(
+    void*, int64_t, double, const int64_t*, int64_t, const char*,
+    double, double, double, double);
+extern "C" int32_t qwen36_set_adapter_id(void*, int64_t, int64_t);
+extern "C" int32_t qwen36_remove_lora(void*, int64_t);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" void* qwen36_get_adapter_optimizer_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t);
+extern "C" void qwen36_free_training_context(void*);
+
+static at::Tensor cuda_rand(std::initializer_list<int64_t> shape) {
+    return at::randn(shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+}
+
+static std::string dp_smoke_sync_dir() {
+    const char* run_id = std::getenv("RUSTRAIN_NCCL_RUN_ID");
+    std::string sanitized = run_id && run_id[0] ? run_id : "native-dp-smoke";
+    for (char& ch : sanitized) {
+        if (!((ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') ||
+              (ch >= '0' && ch <= '9') || ch == '-' || ch == '_')) {
+            ch = '_';
+        }
+    }
+    mkdir("/tmp/rustrain-nccl", 0777);
+    std::string path = "/tmp/rustrain-nccl/" + sanitized;
+    mkdir(path.c_str(), 0777);
+    return path;
+}
+
+static void write_tensor_file(const std::string& path, const at::Tensor& tensor) {
+    auto cpu = tensor.to(at::kCPU).to(at::kFloat).contiguous();
+    const std::string temporary_path = path + ".tmp." +
+        std::to_string(static_cast<long long>(getpid()));
+    FILE* file = std::fopen(temporary_path.c_str(), "wb");
+    assert(file);
+    const int64_t rows = cpu.size(0);
+    const int64_t cols = cpu.size(1);
+    assert(std::fwrite(&rows, sizeof(rows), 1, file) == 1);
+    assert(std::fwrite(&cols, sizeof(cols), 1, file) == 1);
+    assert(std::fwrite(cpu.data_ptr<float>(), sizeof(float), cpu.numel(), file) ==
+        static_cast<size_t>(cpu.numel()));
+    std::fclose(file);
+    assert(std::rename(temporary_path.c_str(), path.c_str()) == 0);
+}
+
+static at::Tensor read_tensor_file(const std::string& path) {
+    FILE* file = nullptr;
+    for (int attempt = 0; attempt < 6000 && !file; ++attempt) {
+        file = std::fopen(path.c_str(), "rb");
+        if (!file) usleep(10000);
+    }
+    assert(file);
+    int64_t rows = 0;
+    int64_t cols = 0;
+    assert(std::fread(&rows, sizeof(rows), 1, file) == 1);
+    assert(std::fread(&cols, sizeof(cols), 1, file) == 1);
+    auto result = at::empty({rows, cols}, at::TensorOptions().dtype(at::kFloat));
+    assert(std::fread(result.data_ptr<float>(), sizeof(float), result.numel(), file) ==
+        static_cast<size_t>(result.numel()));
+    std::fclose(file);
+    return result;
+}
+
+struct DpAdapterInitial {
+    at::Tensor shared_a;
+    at::Tensor shared_b;
+    at::Tensor expert_a;
+    at::Tensor expert_b;
+};
+
+static DpAdapterInitial capture_dp_adapter(void* ctx, int64_t adapter_id) {
+    auto* shared_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_id, 0, "shared_gate_proj", 0));
+    auto* shared_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_id, 0, "shared_gate_proj", 1));
+    auto* expert_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_id, 0, "experts_gate_up_proj", 0));
+    auto* expert_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_id, 0, "experts_gate_up_proj", 1));
+    assert(shared_a && shared_b && expert_a && expert_b);
+    return {
+        shared_a->clone(), shared_b->clone(),
+        expert_a->clone(), expert_b->clone()
+    };
+}
+
+static void restore_dp_adapter(
+    void* ctx, int64_t adapter_id, const DpAdapterInitial& initial
+) {
+    auto shared_a = initial.shared_a;
+    auto shared_b = initial.shared_b;
+    auto expert_a = initial.expert_a;
+    auto expert_b = initial.expert_b;
+    assert(qwen36_set_adapter_lora_tensor(
+        ctx, adapter_id, 0, "shared_gate_proj", 0, &shared_a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        ctx, adapter_id, 0, "shared_gate_proj", 1, &shared_b) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        ctx, adapter_id, 0, "experts_gate_up_proj", 0, &expert_a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        ctx, adapter_id, 0, "experts_gate_up_proj", 1, &expert_b) == 0);
+}
+
+static at::Tensor dp_adapter_b_delta(
+    void* ctx, int64_t adapter_id, const DpAdapterInitial& initial
+) {
+    auto* shared_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_id, 0, "shared_gate_proj", 1));
+    auto* expert_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_id, 0, "experts_gate_up_proj", 1));
+    assert(shared_b && expert_b);
+    return at::cat({
+        (*shared_b - initial.shared_b).to(at::kFloat).reshape({-1}),
+        (*expert_b - initial.expert_b).to(at::kFloat).reshape({-1})
+    });
+}
+
+static at::Tensor dp_adapter_b_optimizer(
+    void* ctx, int64_t adapter_id, bool is_v
+) {
+    auto* shared_state = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_optimizer_tensor(
+            ctx, adapter_id, 0, "shared_gate_proj", 1, is_v ? 1 : 0));
+    auto* expert_state = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_optimizer_tensor(
+            ctx, adapter_id, 0, "experts_gate_up_proj", 1, is_v ? 1 : 0));
+    assert(shared_state && expert_state);
+    assert(shared_state->scalar_type() == at::kFloat);
+    assert(expert_state->scalar_type() == at::kFloat);
+    return at::cat({
+        shared_state->reshape({-1}), expert_state->reshape({-1})});
+}
+
+static at::Tensor dp_adapter_initial_b(const DpAdapterInitial& initial) {
+    return at::cat({
+        initial.shared_b.to(at::kFloat).reshape({-1}),
+        initial.expert_b.to(at::kFloat).reshape({-1})
+    });
+}
+
+static int run_dynamic_dp_smoke(
+    std::vector<void*>& weight_ptrs,
+    at::Tensor& embed, at::Tensor& final_norm, at::Tensor& lm_head,
+    LayerConfig& config, int process_rank, int64_t lora_rank,
+    int64_t vocab, int64_t intermediate, int64_t experts
+) {
+    constexpr double learning_rate = 1e-3;
+    constexpr double adam_eps = 1e-8;
+    constexpr double beta1 = 0.9;
+    constexpr double beta2 = 0.999;
+    const int64_t target_layer = 0;
+    const char* targets = "shared_gate_proj,experts_gate_up_proj";
+    auto create_context = [&]() {
+        return qwen36_create_training_context(
+            weight_ptrs.data(), static_cast<int64_t>(weight_ptrs.size()),
+            &embed, &final_norm, &lm_head, &config, 1,
+            static_cast<int32_t>(at::kBFloat16),
+            1.0, learning_rate, beta1, beta2, adam_eps,
+            vocab, 1e-5, lora_rank, &target_layer, 1, targets);
+    };
+    auto add_adapters = [&](void* ctx) {
+        std::vector<int64_t> ids;
+        ids.push_back(qwen36_add_lora(
+            ctx, lora_rank, 1.0, &target_layer, 1, targets));
+        ids.push_back(qwen36_add_lora(
+            ctx, lora_rank, 1.0, &target_layer, 1, targets));
+        assert(ids[0] > 0 && ids[1] > ids[0]);
+        return ids;
+    };
+
+    auto input_ids = process_rank == 0
+        ? at::tensor({1, 2, 3, 4, 4, 3, 2, 1},
+            at::TensorOptions().device(at::kCUDA).dtype(at::kLong))
+        : at::tensor({2, 4, 1, 3, 3, 1, 4, 2},
+            at::TensorOptions().device(at::kCUDA).dtype(at::kLong));
+    input_ids = input_ids.reshape({2, 4});
+    auto target_mask = process_rank == 0
+        ? at::tensor({1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0},
+            at::TensorOptions().device(at::kCUDA).dtype(at::kFloat))
+        : at::tensor({1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0},
+            at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    target_mask = target_mask.reshape({2, 4});
+    auto attention_mask = at::ones({2, 4},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+
+    // First obtain each rank's independently-normalized local first moment.
+    // At the first step m=(1-beta1)*g, so token weighting remains linear and
+    // can be checked without inferring gradients from BF16 parameter deltas.
+    void* local_ctx = create_context();
+    assert(local_ctx);
+    auto local_ids = add_adapters(local_ctx);
+    std::vector<DpAdapterInitial> initial;
+    initial.push_back(capture_dp_adapter(local_ctx, local_ids[0]));
+    initial.push_back(capture_dp_adapter(local_ctx, local_ids[1]));
+    const double local_loss = qwen36_train_multi_lora(
+        local_ctx, &input_ids, &target_mask, &attention_mask, 2, lora_rank);
+    assert(local_loss == local_loss && local_loss > 0.0);
+    c10::cuda::device_synchronize();
+    auto local_gradients = at::stack({
+        dp_adapter_b_optimizer(local_ctx, local_ids[0], false),
+        dp_adapter_b_optimizer(local_ctx, local_ids[1], false)
+    }).to(at::kCPU);
+    qwen36_free_training_context(local_ctx);
+
+    const std::string sync_dir = dp_smoke_sync_dir();
+    const std::string local_path = sync_dir + "/dp-local-" +
+        std::to_string(process_rank) + ".bin";
+    const std::string peer_path = sync_dir + "/dp-local-" +
+        std::to_string(1 - process_rank) + ".bin";
+    std::remove(local_path.c_str());
+    write_tensor_file(local_path, local_gradients);
+    auto peer_gradients = read_tensor_file(peer_path);
+    assert(peer_gradients.sizes() == local_gradients.sizes());
+
+    void* distributed_ctx = create_context();
+    assert(distributed_ctx);
+    assert(qwen36_init_nccl(distributed_ctx) == 0);
+    auto distributed_ids = add_adapters(distributed_ctx);
+    restore_dp_adapter(distributed_ctx, distributed_ids[0], initial[0]);
+    restore_dp_adapter(distributed_ctx, distributed_ids[1], initial[1]);
+    const double distributed_loss = qwen36_train_multi_lora(
+        distributed_ctx, &input_ids, &target_mask, &attention_mask,
+        2, lora_rank);
+    assert(distributed_loss == distributed_loss && distributed_loss > 0.0);
+    c10::cuda::device_synchronize();
+    auto distributed_deltas = at::stack({
+        dp_adapter_b_delta(
+            distributed_ctx, distributed_ids[0], initial[0]),
+        dp_adapter_b_delta(
+            distributed_ctx, distributed_ids[1], initial[1])
+    }).to(at::kCPU);
+    auto distributed_gradients = at::stack({
+        dp_adapter_b_optimizer(distributed_ctx, distributed_ids[0], false),
+        dp_adapter_b_optimizer(distributed_ctx, distributed_ids[1], false)
+    }).to(at::kCPU);
+    auto distributed_v = at::stack({
+        dp_adapter_b_optimizer(distributed_ctx, distributed_ids[0], true),
+        dp_adapter_b_optimizer(distributed_ctx, distributed_ids[1], true)
+    }).to(at::kCPU);
+    auto initial_b = at::stack({
+        dp_adapter_initial_b(initial[0]), dp_adapter_initial_b(initial[1])
+    }).to(at::kCPU);
+
+    const std::string distributed_path = sync_dir + "/dp-global-" +
+        std::to_string(process_rank) + ".bin";
+    const std::string peer_distributed_path = sync_dir + "/dp-global-" +
+        std::to_string(1 - process_rank) + ".bin";
+    std::remove(distributed_path.c_str());
+    write_tensor_file(distributed_path, distributed_gradients);
+    auto peer_distributed_gradients = read_tensor_file(peer_distributed_path);
+    assert(at::allclose(
+        distributed_gradients, peer_distributed_gradients, 0.0, 0.0));
+
+    auto gradient0 = process_rank == 0 ? local_gradients : peer_gradients;
+    auto gradient1 = process_rank == 0 ? peer_gradients : local_gradients;
+    auto counts0 = at::tensor(
+        {1.0, 3.0}, at::TensorOptions().dtype(at::kFloat)).reshape({2, 1});
+    auto counts1 = at::tensor(
+        {3.0, 1.0}, at::TensorOptions().dtype(at::kFloat)).reshape({2, 1});
+    auto expected_gradient =
+        (gradient0 * counts0 + gradient1 * counts1) / (counts0 + counts1);
+    auto aggregate_count_gradient = (gradient0 + gradient1) * 0.5;
+
+    auto relative_error = [&](const at::Tensor& actual,
+                              const at::Tensor& expected) {
+        return (actual - expected).abs().sum().item<double>() /
+            std::max(expected.abs().sum().item<double>(), 1e-12);
+    };
+    const int64_t shared_numel = intermediate * lora_rank;
+    const int64_t expert_numel = experts * 2 * intermediate * lora_rank;
+    assert(distributed_gradients.size(1) == shared_numel + expert_numel);
+    const double all_error = relative_error(
+        distributed_gradients, expected_gradient);
+    const double grouped_error = relative_error(
+        distributed_gradients.narrow(1, shared_numel, expert_numel),
+        expected_gradient.narrow(1, shared_numel, expert_numel));
+    const double old_formula_gap =
+        (expected_gradient - aggregate_count_gradient).abs().sum().item<double>();
+    auto expected_v = distributed_gradients.square() *
+        ((1.0 - beta2) / ((1.0 - beta1) * (1.0 - beta1)));
+    const double v_error = relative_error(distributed_v, expected_v);
+    auto expected_parameter = (
+        initial_b - learning_rate *
+            (distributed_gradients / (1.0 - beta1)) /
+            ((distributed_v / (1.0 - beta2)).sqrt() + adam_eps))
+        .to(at::kBFloat16).to(at::kFloat);
+    auto expected_parameter_delta = expected_parameter - initial_b;
+    const double adam_delta_max_error =
+        (distributed_deltas - expected_parameter_delta)
+            .abs().max().item<double>();
+    const double grouped_update = distributed_deltas
+        .narrow(1, shared_numel, expert_numel).abs().sum().item<double>();
+    std::printf(
+        "native_qwen36_dynamic_dp_weighting rank=%d loss=%0.8f "
+        "relative_error=%0.6e grouped_error=%0.6e old_formula_gap=%0.6e "
+        "v_error=%0.6e adam_delta_max_error=%0.6e grouped_update=%0.6e\n",
+        process_rank, distributed_loss, all_error, grouped_error,
+        old_formula_gap, v_error, adam_delta_max_error, grouped_update);
+    assert(old_formula_gap > 1e-8);
+    assert(grouped_update > 0.0);
+    assert(all_error < 2e-5);
+    assert(grouped_error < 2e-5);
+    assert(v_error < 2e-5);
+    assert(adam_delta_max_error < 2e-5);
+    qwen36_free_training_context(distributed_ctx);
+    return 0;
+}
+
+int main() {
+    assert(qwen36_kernel_abi_version() == 31);
+    const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1");
+    const int process_rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0");
+    const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0");
+    assert(world == 1 || (world == 2 && process_rank >= 0 && process_rank < world));
+    const int tp_size = std::atoi(std::getenv("TP_SIZE") ? std::getenv("TP_SIZE") : "1");
+    assert(tp_size == 1 || (tp_size == world && tp_size == 2));
+    c10::cuda::CUDAGuard guard(local_rank);
+    at::manual_seed(7);
+
+    constexpr int64_t hidden = 16;
+    constexpr int64_t vocab = 8;
+    constexpr int64_t experts = 2;
+    constexpr int64_t head_dim = 8;
+    constexpr int64_t intermediate = 8;
+    constexpr int64_t rank = 8;
+    const int64_t local_lora_rank = rank / tp_size;
+
+    // One full-attention MoE layer. Shapes intentionally match the native
+    // weight order used by build_weight_ptrs/kernel.cpp.
+    std::vector<at::Tensor> weights;
+    weights.push_back(cuda_rand({hidden}));                 // input RMSNorm
+    weights.push_back(cuda_rand({hidden}));                 // post-attention RMSNorm
+    weights.push_back(cuda_rand({2 * head_dim, hidden}));   // q_proj
+    weights.push_back(cuda_rand({head_dim}));               // q_norm
+    weights.push_back(cuda_rand({head_dim, hidden}));       // k_proj
+    weights.push_back(cuda_rand({head_dim}));               // k_norm
+    weights.push_back(cuda_rand({head_dim, hidden}));       // v_proj
+    weights.push_back(cuda_rand({hidden, head_dim}));       // o_proj
+    weights.push_back(cuda_rand({experts, hidden}));        // router
+    weights.push_back(cuda_rand({1, hidden}));              // shared expert gate
+    weights.push_back(cuda_rand({intermediate, hidden}));   // shared gate
+    weights.push_back(cuda_rand({intermediate, hidden}));   // shared up
+    weights.push_back(cuda_rand({hidden, intermediate}));   // shared down
+    weights.push_back(cuda_rand({experts, 2 * intermediate, hidden}));
+    weights.push_back(cuda_rand({experts, hidden, intermediate}));
+    for (auto& weight : weights) weight.set_requires_grad(false);
+
+    auto embed = cuda_rand({vocab, hidden});
+    auto final_norm = cuda_rand({hidden});
+    auto lm_head = cuda_rand({vocab, hidden});
+    embed.set_requires_grad(false);
+    final_norm.set_requires_grad(false);
+    lm_head.set_requires_grad(false);
+
+    std::vector<void*> weight_ptrs;
+    weight_ptrs.reserve(weights.size());
+    for (auto& weight : weights) weight_ptrs.push_back(&weight);
+    LayerConfig config{};
+    config.layer_type = 0;
+    config.num_heads = 1;
+    config.num_kv_heads = 1;
+    config.head_dim = head_dim;
+    config.num_k_heads = 0;
+    config.key_dim = 0;
+    config.num_v_heads = 0;
+    config.val_dim = 0;
+    config.conv_kernel = 0;
+    config.partial_rotary_factor = 1.0;
+    config.rope_theta = 10000.0;
+    config.rms_eps = 1e-5;
+    config.num_experts = experts;
+    config.top_k = 1;
+    config.moe_intermediate = intermediate;
+    config.expert_start = 0;
+    config.expert_count = experts;
+    config.intermediate_size = 0;
+    config.norm_topk_prob = 1;
+    config.nccl_comm = nullptr;
+    config.nccl_stream = nullptr;
+
+    const bool data_parallel = std::getenv("RUSTRAIN_DATA_PARALLEL") &&
+        std::strcmp(std::getenv("RUSTRAIN_DATA_PARALLEL"), "0") != 0;
+    if (world == 2 && tp_size == 1 && data_parallel) {
+        return run_dynamic_dp_smoke(
+            weight_ptrs, embed, final_norm, lm_head, config,
+            process_rank, rank, vocab, intermediate, experts);
+    }
+
+    const int64_t target_layer = 0;
+    // Native C++ must reject a mixed TP/DP/EP topology even when Rust-side
+    // validation is bypassed. The smoke process is single-rank, so simulate
+    // an invalid world before creating the real context.
+    if (world == 1 && tp_size == 1) {
+        setenv("TP_SIZE", "2", 1);
+        setenv("WORLD_SIZE", "4", 1);
+        void* invalid_topology_ctx = qwen36_create_training_context(
+            weight_ptrs.data(), static_cast<int64_t>(weight_ptrs.size()),
+            &embed, &final_norm, &lm_head, &config, 1,
+            static_cast<int32_t>(at::kBFloat16),
+            1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank,
+            &target_layer, 1, "experts_gate_up_proj,experts_down_proj");
+        assert(invalid_topology_ctx == nullptr);
+        setenv("TP_SIZE", "1", 1);
+        setenv("WORLD_SIZE", "1", 1);
+    }
+    if (world == 1 && tp_size == 1) {
+        setenv("PP_SIZE", "2", 1);
+        void* invalid_pp_ctx = qwen36_create_training_context(
+            weight_ptrs.data(), static_cast<int64_t>(weight_ptrs.size()),
+            &embed, &final_norm, &lm_head, &config, 1,
+            static_cast<int32_t>(at::kBFloat16),
+            1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank,
+            &target_layer, 1, "experts_gate_up_proj,experts_down_proj");
+        assert(invalid_pp_ctx == nullptr);
+        unsetenv("PP_SIZE");
+    }
+    void* ctx = qwen36_create_training_context(
+        weight_ptrs.data(), static_cast<int64_t>(weight_ptrs.size()),
+        &embed, &final_norm, &lm_head, &config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank,
+        &target_layer, 1, "experts_gate_up_proj,experts_down_proj");
+    if (!ctx) return 2;
+    if (world > 1) assert(qwen36_init_nccl(ctx) == 0);
+
+    const int64_t count = qwen36_get_lora_count(ctx);
+    assert(count == 9);
+    auto* expert_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(ctx, 7));
+    auto* expert_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(ctx, 7));
+    auto* down_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(ctx, 8));
+    auto* down_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(ctx, 8));
+    assert(expert_a && expert_b && down_a && down_b);
+    assert(expert_a->sizes() == at::IntArrayRef({experts, local_lora_rank, hidden}));
+    assert(expert_b->sizes() == at::IntArrayRef({experts, 2 * intermediate, local_lora_rank}));
+    assert(down_a->sizes() == at::IntArrayRef({experts, local_lora_rank, intermediate}));
+    assert(down_b->sizes() == at::IntArrayRef({experts, hidden, local_lora_rank}));
+
+    // Make both B tensors nonzero so the step exercises the LoRA branches.
+    auto expert_b_value = at::ones(expert_b->sizes(), expert_b->options());
+    auto down_b_value = at::ones(down_b->sizes(), down_b->options());
+    assert(qwen36_set_lora_tensor(ctx, 7, 1, &expert_b_value) == 0);
+    assert(qwen36_set_lora_tensor(ctx, 8, 1, &down_b_value) == 0);
+    auto expert_a_before = expert_a->clone();
+    auto expert_b_before = expert_b->clone();
+
+    auto input_ids = at::arange(1, 3, at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 2});
+    auto target_mask = at::ones({1, 2}, at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    auto attention_mask = at::ones({1, 2}, at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+    // Distinct tenant rows exercise the production [n_total, seq] path. The
+    // rows are intentionally different so a repeated batch-1 implementation
+    // cannot satisfy the assertions below.
+    auto multi_input_ids = at::tensor({1, 2, 3, 4},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({2, 2});
+    auto multi_target_mask = at::ones({2, 2},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    auto multi_attention_mask = at::ones({2, 2},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+
+    // The optimized and fallback routed-expert paths must agree before the
+    // optimizer changes any parameters. On older libtorch both calls use the
+    // fallback, preserving the same compatibility check.
+    setenv("QWEN36_DISABLE_GROUPED_MM", "1", 1);
+    const double fallback_loss =
+        qwen36_eval_step(ctx, &input_ids, &target_mask, &attention_mask);
+    unsetenv("QWEN36_DISABLE_GROUPED_MM");
+    setenv("QWEN36_REPORT_GROUPED_MM", "1", 1);
+    const double grouped_loss =
+        qwen36_eval_step(ctx, &input_ids, &target_mask, &attention_mask);
+    const int64_t host_input_ids[] = {1, 2};
+    const int64_t host_target_mask[] = {1, 1};
+    const int64_t host_attention_mask[] = {1, 1};
+    const double host_eval_loss = qwen36_eval_step_host_i64(
+        ctx, host_input_ids, host_target_mask, host_attention_mask, 1, 2);
+    std::printf(
+        "native_qwen36_moe_lora_parity fallback=%0.8f grouped=%0.8f diff=%0.8e\n",
+        fallback_loss, grouped_loss, std::abs(fallback_loss - grouped_loss));
+    assert(fallback_loss > 0.0 && grouped_loss > 0.0 && host_eval_loss > 0.0);
+    assert(std::abs(fallback_loss - grouped_loss) <= 2e-2);
+    assert(std::abs(grouped_loss - host_eval_loss) <= 2e-2);
+
+    if (world == 1) {
+        auto make_router_context = [&]() {
+            return qwen36_create_training_context(
+                weight_ptrs.data(), static_cast<int64_t>(weight_ptrs.size()),
+                &embed, &final_norm, &lm_head, &config, 1,
+                static_cast<int32_t>(at::kBFloat16),
+                1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank,
+                &target_layer, 1, "q_proj");
+        };
+        at::manual_seed(1234);
+        void* router_reference = make_router_context();
+        at::manual_seed(1234);
+        void* router_aux = make_router_context();
+        assert(router_reference && router_aux);
+        assert(qwen36_set_router_aux_loss_coef(router_aux, -1.0) != 0);
+        assert(qwen36_set_router_aux_loss_coef(router_aux, 0.01) == 0);
+        auto* reference_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(router_reference, 0));
+        auto* aux_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(router_aux, 0));
+        assert(reference_b && aux_b && reference_b->sizes() == aux_b->sizes());
+        auto nonzero_b = at::full(reference_b->sizes(), 0.01, reference_b->options());
+        assert(qwen36_set_lora_tensor(router_reference, 0, 1, &nonzero_b) == 0);
+        assert(qwen36_set_lora_tensor(router_aux, 0, 1, &nonzero_b) == 0);
+        qwen36_set_checkpoint(router_reference, 1, 1);
+        qwen36_set_checkpoint(router_aux, 1, 1);
+        const double reference_loss = qwen36_train_step(
+            router_reference, &input_ids, &target_mask, &attention_mask);
+        const double aux_loss = qwen36_train_step(
+            router_aux, &input_ids, &target_mask, &attention_mask);
+        std::vector<void*> reference_m(18), reference_v(18), aux_m(18), aux_v(18);
+        assert(qwen36_export_optimizer_state(
+            router_reference, reference_m.data(), reference_v.data(), 18) == 18);
+        assert(qwen36_export_optimizer_state(
+            router_aux, aux_m.data(), aux_v.data(), 18) == 18);
+        auto* reference_a_m = reinterpret_cast<at::Tensor*>(reference_m[0]);
+        auto* aux_a_m = reinterpret_cast<at::Tensor*>(aux_m[0]);
+        assert(reference_a_m && aux_a_m);
+        const double router_m_gap =
+            (*reference_a_m - *aux_a_m).abs().sum().item<double>();
+        std::printf(
+            "native_qwen36_router_aux reference_loss=%0.8f aux_loss=%0.8f "
+            "loss_gap=%0.8e first_m_gap=%0.8e\n",
+            reference_loss, aux_loss, aux_loss - reference_loss, router_m_gap);
+        assert(aux_loss > reference_loss);
+        assert(router_m_gap > 0.0);
+        assert(qwen36_add_lora(
+            router_aux, rank, 1.0, &target_layer, 1, "q_proj") < 0);
+        qwen36_free_training_context(router_aux);
+        qwen36_free_training_context(router_reference);
+    }
+
+    const double loss = qwen36_train_step_host_i64(
+        ctx, host_input_ids, host_target_mask, host_attention_mask, 1, 2);
+    c10::cuda::device_synchronize();
+    std::printf("native_qwen36_moe_lora_smoke loss=%0.8f\n", loss);
+    assert(loss == loss);
+    assert(loss > 0.0);
+    const double update_norm = (*expert_a - expert_a_before).abs().sum().item<double>();
+    const double b_update_norm = (*expert_b - expert_b_before).abs().sum().item<double>();
+    std::vector<void*> fixed_m(2 * count);
+    std::vector<void*> fixed_v(2 * count);
+    assert(qwen36_export_optimizer_state(
+        ctx, fixed_m.data(), fixed_v.data(), 2 * count) == 2 * count);
+    auto* expert_b_m = reinterpret_cast<at::Tensor*>(fixed_m[2 * 7 + 1]);
+    auto* expert_b_v = reinterpret_cast<at::Tensor*>(fixed_v[2 * 7 + 1]);
+    assert(expert_b_m && expert_b_v);
+    auto expected_expert_b = (
+        expert_b_before.to(at::kFloat) - 1e-3 *
+            (expert_b_m->to(at::kFloat) / (1.0 - 0.9)) /
+            ((expert_b_v->to(at::kFloat) / (1.0 - 0.999)).sqrt() + 1e-8))
+        .to(at::kBFloat16);
+    const double fixed_adam_max_error =
+        (*expert_b - expected_expert_b).abs().max().item<double>();
+    std::printf("native_qwen36_moe_lora_smoke expert_a_update=%0.8e expert_b_update=%0.8e\n",
+                update_norm, b_update_norm);
+    assert(update_norm > 0.0 || b_update_norm > 0.0);
+    std::printf(
+        "native_qwen36_fixed_adam_oracle max_error=%0.8e\n",
+        fixed_adam_max_error);
+    assert(fixed_adam_max_error == 0.0);
+
+    // Dynamic multi-adapter batches must apply shared-expert MLP LoRA per
+    // sample and preserve the parameter updates after chunk registry restore.
+    const char* shared_targets =
+        "shared_gate_proj,shared_up_proj,shared_down_proj,"
+        "experts_gate_up_proj,experts_down_proj";
+    assert(qwen36_add_lora(
+        ctx, rank, INFINITY, &target_layer, 1, shared_targets) < 0);
+    const int64_t adapter_one = qwen36_add_lora(
+        ctx, rank, 1.0, &target_layer, 1, shared_targets);
+    const int64_t adapter_two = qwen36_add_lora(
+        ctx, rank, 1.0, &target_layer, 1, shared_targets);
+    assert(adapter_one == 1 && adapter_two == 2);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 0);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 0);
+    assert(qwen36_set_adapter_step_count(ctx, adapter_one, -1) != 0);
+    assert(qwen36_set_adapter_step_count(ctx, adapter_one, 4) == 0);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 4);
+    assert(qwen36_set_adapter_step_count(ctx, adapter_one, 0) == 0);
+    const int64_t initial_adapter_ids[] = {adapter_one, adapter_two};
+    const int64_t initial_adapter_steps[] = {0, 0};
+    const int64_t stale_adapter_steps[] = {1, 0};
+    assert(qwen36_validate_adapter_steps_v1(
+        ctx, initial_adapter_ids, initial_adapter_steps, 2) == 0);
+    assert(qwen36_validate_adapter_steps_v1(
+        ctx, initial_adapter_ids, stale_adapter_steps, 2) != 0);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 0);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 0);
+    auto* dynamic_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_one, 0, "shared_gate_proj", 1));
+    assert(dynamic_b && dynamic_b->sizes() == at::IntArrayRef({intermediate, local_lora_rank}));
+    auto dynamic_b_value = at::full(
+        dynamic_b->sizes(), 0.01, dynamic_b->options());
+    assert(qwen36_set_adapter_lora_tensor(
+        ctx, adapter_one, 0, "shared_gate_proj", 1, &dynamic_b_value) == 0);
+    auto dynamic_b_before = dynamic_b->clone();
+    auto* dynamic_b_two = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_two, 0, "shared_gate_proj", 1));
+    assert(dynamic_b_two && dynamic_b_two->sizes() == dynamic_b->sizes());
+    auto dynamic_b_two_value = at::full(
+        dynamic_b_two->sizes(), -0.01, dynamic_b_two->options());
+    assert(qwen36_set_adapter_lora_tensor(
+        ctx, adapter_two, 0, "shared_gate_proj", 1, &dynamic_b_two_value) == 0);
+    auto dynamic_b_two_before = dynamic_b_two->clone();
+    auto* dynamic_expert_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_one, 0, "experts_gate_up_proj", 1));
+    assert(dynamic_expert_b && dynamic_expert_b->sizes() ==
+        at::IntArrayRef({experts, 2 * intermediate, local_lora_rank}));
+    auto dynamic_expert_b_value = at::full(
+        dynamic_expert_b->sizes(), 0.01, dynamic_expert_b->options());
+    assert(qwen36_set_adapter_lora_tensor(
+        ctx, adapter_one, 0, "experts_gate_up_proj", 1,
+        &dynamic_expert_b_value) == 0);
+    auto dynamic_expert_b_before = dynamic_expert_b->clone();
+    const double multi_loss = qwen36_train_multi_lora(
+        ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask, 2, rank);
+    c10::cuda::device_synchronize();
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 1);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1);
+    dynamic_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_one, 0, "shared_gate_proj", 1));
+    assert(dynamic_b);
+    dynamic_expert_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_one, 0, "experts_gate_up_proj", 1));
+    assert(dynamic_expert_b);
+    const double dynamic_update =
+        (*dynamic_b - dynamic_b_before).abs().sum().item<double>();
+    const double dynamic_expert_update =
+        (*dynamic_expert_b - dynamic_expert_b_before).abs().sum().item<double>();
+    dynamic_b_two = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_two, 0, "shared_gate_proj", 1));
+    assert(dynamic_b_two);
+    const double dynamic_two_update =
+        (*dynamic_b_two - dynamic_b_two_before).abs().sum().item<double>();
+    std::printf(
+        "native_qwen36_multi_lora_smoke loss=%0.8f shared_gate_b_update=%0.8e "
+        "expert_gate_up_b_update=%0.8e adapter_two_gate_b_update=%0.8e\n",
+        multi_loss, dynamic_update, dynamic_expert_update, dynamic_two_update);
+    assert(multi_loss == multi_loss && multi_loss > 0.0);
+    assert(dynamic_update > 0.0);
+    assert(dynamic_expert_update > 0.0);
+    assert(dynamic_two_update > 0.0);
+
+    double selected_eval_losses[2] = {-1.0, -1.0};
+    const int64_t eval_adapter_ids[] = {adapter_one, adapter_two};
+    assert(qwen36_eval_multi_lora_host_i64_v1(
+        ctx, host_input_ids, host_target_mask, host_attention_mask,
+        1, 2, eval_adapter_ids, 2, selected_eval_losses, 2) == 0);
+    std::printf(
+        "native_qwen36_selected_multi_lora_eval adapter_one=%0.8f "
+        "adapter_two=%0.8f\n",
+        selected_eval_losses[0], selected_eval_losses[1]);
+    assert(selected_eval_losses[0] >= 0.0 && selected_eval_losses[1] >= 0.0);
+
+    auto adapter_one_before_selected = dynamic_b->clone();
+    auto adapter_two_before_selected = dynamic_b_two->clone();
+    auto selected_input_ids = multi_input_ids.narrow(0, 0, 1).contiguous();
+    auto selected_target_mask = multi_target_mask.narrow(0, 0, 1).contiguous();
+    auto selected_attention_mask = multi_attention_mask.narrow(0, 0, 1).contiguous();
+    const int64_t selected_adapter_ids[] = {adapter_one};
+    const double selected_loss = qwen36_train_multi_lora_host_i64(
+        ctx, host_input_ids, host_target_mask, host_attention_mask,
+        1, 2, 1, rank, selected_adapter_ids, 1);
+    c10::cuda::device_synchronize();
+    assert(selected_loss == selected_loss && selected_loss > 0.0);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 2);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1);
+    dynamic_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_one, 0, "shared_gate_proj", 1));
+    dynamic_b_two = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, adapter_two, 0, "shared_gate_proj", 1));
+    assert(dynamic_b && dynamic_b_two);
+    const double selected_update =
+        (*dynamic_b - adapter_one_before_selected).abs().sum().item<double>();
+    const double unselected_update =
+        (*dynamic_b_two - adapter_two_before_selected).abs().sum().item<double>();
+    std::printf(
+        "native_qwen36_selected_multi_lora_smoke loss=%0.8f "
+        "selected_update=%0.8e unselected_update=%0.8e\n",
+        selected_loss, selected_update, unselected_update);
+    assert(selected_update > 0.0);
+    assert(unselected_update == 0.0);
+
+    auto* transactional_m = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_optimizer_tensor(
+            ctx, adapter_one, 0, "shared_gate_proj", 1, 0));
+    auto* transactional_v = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_optimizer_tensor(
+            ctx, adapter_one, 0, "shared_gate_proj", 1, 1));
+    assert(transactional_m && transactional_v);
+    auto transactional_b_before = dynamic_b->clone();
+    auto transactional_m_before = transactional_m->clone();
+    auto transactional_v_before = transactional_v->clone();
+    setenv("QWEN36_TEST_FAIL_DYNAMIC_ADAM_BEFORE_COMMIT", "1", 1);
+    const double injected_failure = qwen36_train_multi_lora_selected(
+        ctx, &selected_input_ids, &selected_target_mask,
+        &selected_attention_mask, selected_adapter_ids, 1, rank);
+    unsetenv("QWEN36_TEST_FAIL_DYNAMIC_ADAM_BEFORE_COMMIT");
+    c10::cuda::device_synchronize();
+    assert(injected_failure < 0.0);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 2);
+    assert((*dynamic_b - transactional_b_before).abs().max().item<double>() == 0.0);
+    assert((*transactional_m - transactional_m_before).abs().max().item<double>() == 0.0);
+    assert((*transactional_v - transactional_v_before).abs().max().item<double>() == 0.0);
+    assert(qwen36_get_context_health(ctx) == 0);
+    std::printf("native_qwen36_transactional_adam_failure_smoke ok\n");
+
+    const int64_t launches_before_non_finite =
+        qwen36_get_dynamic_adam_launch_count(ctx);
+    setenv("QWEN36_TEST_INJECT_DYNAMIC_GRAD_NAN", "1", 1);
+    const double non_finite_gradient_failure =
+        qwen36_train_multi_lora_selected(
+            ctx, &selected_input_ids, &selected_target_mask,
+            &selected_attention_mask, selected_adapter_ids, 1, rank);
+    unsetenv("QWEN36_TEST_INJECT_DYNAMIC_GRAD_NAN");
+    c10::cuda::device_synchronize();
+    assert(non_finite_gradient_failure < 0.0);
+    assert(qwen36_get_dynamic_adam_launch_count(ctx) ==
+        launches_before_non_finite);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 2);
+    assert((*dynamic_b - transactional_b_before).abs().max().item<double>() == 0.0);
+    assert((*transactional_m - transactional_m_before).abs().max().item<double>() == 0.0);
+    assert((*transactional_v - transactional_v_before).abs().max().item<double>() == 0.0);
+    assert(qwen36_get_accumulation_active(ctx) == 0);
+    assert(qwen36_get_accumulated_token_weight(ctx) == 0.0);
+    assert(qwen36_get_context_health(ctx) == 0);
+    std::printf("native_qwen36_dynamic_non_finite_gradient_smoke ok\n");
+
+    const int64_t unknown_adapter_ids[] = {adapter_two + 1000};
+    assert(qwen36_train_multi_lora_selected(
+        ctx, &selected_input_ids, &selected_target_mask,
+        &selected_attention_mask, unknown_adapter_ids, 1, rank) < 0.0);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 2);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1);
+    assert(qwen36_get_adapter_lora_tensor(
+        ctx, adapter_one, 0, "shared_gate_proj", 1) != nullptr);
+    assert(qwen36_get_adapter_lora_tensor(
+        ctx, adapter_two, 0, "shared_gate_proj", 1) != nullptr);
+    // A globally empty tenant is a successful no-op: active tenants still
+    // update, while the empty tenant preserves parameters, optimizer state,
+    // and its private Adam clock.
+    auto zero_target_mask = multi_target_mask.clone();
+    zero_target_mask.select(0, 1).zero_();
+    auto empty_tenant_b_before = dynamic_b_two->clone();
+    auto* empty_tenant_m = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_optimizer_tensor(
+            ctx, adapter_two, 0, "shared_gate_proj", 1, 0));
+    auto* empty_tenant_v = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_optimizer_tensor(
+            ctx, adapter_two, 0, "shared_gate_proj", 1, 1));
+    assert(empty_tenant_m && empty_tenant_v);
+    auto empty_tenant_m_before = empty_tenant_m->clone();
+    auto empty_tenant_v_before = empty_tenant_v->clone();
+    assert(qwen36_train_multi_lora(
+        ctx, &multi_input_ids, &zero_target_mask, &multi_attention_mask,
+        2, rank) > 0.0);
+    c10::cuda::device_synchronize();
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 3);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1);
+    assert((*dynamic_b_two - empty_tenant_b_before).abs().max().item<double>() == 0.0);
+    assert((*empty_tenant_m - empty_tenant_m_before).abs().max().item<double>() == 0.0);
+    assert((*empty_tenant_v - empty_tenant_v_before).abs().max().item<double>() == 0.0);
+    assert(qwen36_add_lora(
+        ctx, rank + 1, 12.0, &target_layer, 1, "q_proj") < 0);
+    const int64_t heterogeneous_restore = qwen36_add_lora_v2(
+        ctx, rank + 1, 12.0, &target_layer, 1, "q_proj");
+    assert(heterogeneous_restore == adapter_two + 1);
+    auto* heterogeneous_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, heterogeneous_restore, 0, "q_proj", 1));
+    assert(heterogeneous_b);
+    assert(qwen36_train_multi_lora(
+        ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask,
+        3, rank) < 0.0);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 3);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1);
+    const int64_t heterogeneous_ids[] = {
+        adapter_one, adapter_two, heterogeneous_restore};
+    auto heterogeneous_b_before = heterogeneous_b->clone();
+    auto one_before_v2 = dynamic_b->clone();
+    auto shared_input = multi_input_ids.narrow(0, 0, 1).contiguous();
+    auto shared_target_row = multi_target_mask.narrow(0, 0, 1).contiguous();
+    auto shared_attention = multi_attention_mask.narrow(0, 0, 1).contiguous();
+    const int64_t duplicate_ids[] = {adapter_one, adapter_one};
+    assert(qwen36_train_multi_lora_selected_v2(
+        ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask,
+        duplicate_ids, 2) < 0.0);
+    const int64_t unknown_ids[] = {adapter_one, heterogeneous_restore + 1000};
+    assert(qwen36_train_multi_lora_selected_v2(
+        ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask,
+        unknown_ids, 2) < 0.0);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 3);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1);
+    assert(qwen36_get_adapter_step_count(ctx, heterogeneous_restore) == 0);
+    const int64_t finalizers_before_v2 =
+        qwen36_get_dynamic_finalizer_count(ctx);
+    const int64_t adam_launches_before_v2 =
+        qwen36_get_dynamic_adam_launch_count(ctx);
+    const int64_t train_batches_before_v2 =
+        qwen36_get_dynamic_train_batch_count(ctx);
+    const double heterogeneous_loss = qwen36_train_multi_lora_selected_v2(
+        ctx, &shared_input, &shared_target_row, &shared_attention,
+        heterogeneous_ids, 3);
+    c10::cuda::device_synchronize();
+    assert(heterogeneous_loss == heterogeneous_loss && heterogeneous_loss > 0.0);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 4);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 2);
+    assert(qwen36_get_adapter_step_count(ctx, heterogeneous_restore) == 1);
+    assert(qwen36_get_dynamic_finalizer_count(ctx) ==
+        finalizers_before_v2 + 1);
+    assert(qwen36_get_dynamic_adam_launch_count(ctx) ==
+        adam_launches_before_v2 + 1);
+    assert(qwen36_get_dynamic_train_batch_count(ctx) ==
+        train_batches_before_v2 + 1);
+    assert((*dynamic_b - one_before_v2).abs().sum().item<double>() > 0.0);
+    assert((*heterogeneous_b - heterogeneous_b_before).abs().sum().item<double>() > 0.0);
+
+    auto rollback_one_b = dynamic_b->clone();
+    auto rollback_heterogeneous_b = heterogeneous_b->clone();
+    auto* rollback_one_m = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_optimizer_tensor(
+            ctx, adapter_one, 0, "shared_gate_proj", 1, 0));
+    assert(rollback_one_m);
+    auto rollback_one_m_before = rollback_one_m->clone();
+    setenv("QWEN36_TEST_FAIL_HETERO_GROUP_AFTER", "1", 1);
+    assert(qwen36_train_multi_lora_selected_v2(
+        ctx, &shared_input, &shared_target_row, &shared_attention,
+        heterogeneous_ids, 3) < 0.0);
+    unsetenv("QWEN36_TEST_FAIL_HETERO_GROUP_AFTER");
+    c10::cuda::device_synchronize();
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 4);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 2);
+    assert(qwen36_get_adapter_step_count(ctx, heterogeneous_restore) == 1);
+    assert((*dynamic_b - rollback_one_b).abs().max().item<double>() == 0.0);
+    assert((*rollback_one_m - rollback_one_m_before).abs().max().item<double>() == 0.0);
+    assert((*heterogeneous_b - rollback_heterogeneous_b).abs().max().item<double>() == 0.0);
+    assert(qwen36_get_dynamic_finalizer_count(ctx) ==
+        finalizers_before_v2 + 1);
+    assert(qwen36_get_dynamic_adam_launch_count(ctx) ==
+        adam_launches_before_v2 + 1);
+
+    const int64_t grouped_batches_before =
+        qwen36_get_dynamic_train_batch_count(ctx);
+    setenv("QWEN36_HETERO_PADDED_BATCH", "0", 1);
+    assert(qwen36_train_multi_lora_host_i64(
+        ctx, host_input_ids, host_target_mask, host_attention_mask,
+        1, 2, 3, rank, nullptr, 0) > 0.0);
+    unsetenv("QWEN36_HETERO_PADDED_BATCH");
+    c10::cuda::device_synchronize();
+    assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 5);
+    assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 3);
+    assert(qwen36_get_adapter_step_count(ctx, heterogeneous_restore) == 2);
+    assert(qwen36_get_dynamic_finalizer_count(ctx) ==
+        finalizers_before_v2 + 2);
+    assert(qwen36_get_dynamic_adam_launch_count(ctx) ==
+        adam_launches_before_v2 + 2);
+    assert(qwen36_get_dynamic_train_batch_count(ctx) ==
+        grouped_batches_before + 2);
+    std::printf("native_qwen36_heterogeneous_v2_smoke loss=%0.8f ok\n",
+        heterogeneous_loss);
+    assert(qwen36_remove_lora(ctx, heterogeneous_restore) == 1);
+    setenv("QWEN36_TEST_POISON_DYNAMIC_RECOVERY", "1", 1);
+    assert(qwen36_train_multi_lora_selected_v2(
+        ctx, &selected_input_ids, &selected_target_mask,
+        &selected_attention_mask, unknown_adapter_ids, 1) < 0.0);
+    unsetenv("QWEN36_TEST_POISON_DYNAMIC_RECOVERY");
+    assert(qwen36_get_context_health(ctx) == -1);
+    assert(qwen36_train_multi_lora_selected_v2(
+        ctx, &selected_input_ids, &selected_target_mask,
+        &selected_attention_mask, selected_adapter_ids, 1) < 0.0);
+    std::printf("native_qwen36_dynamic_recovery_poison_smoke ok\n");
+    qwen36_free_training_context(ctx);
+
+    // Dense Qwen3.5 variants use the same per-sample activation path for
+    // gate/up/down projections.
+    std::vector<at::Tensor> dense_weights(weights.begin(), weights.begin() + 8);
+    dense_weights.push_back(cuda_rand({intermediate, hidden}));
+    dense_weights.push_back(cuda_rand({intermediate, hidden}));
+    dense_weights.push_back(cuda_rand({hidden, intermediate}));
+    weight_ptrs.clear();
+    for (auto& weight : dense_weights) weight_ptrs.push_back(&weight);
+    LayerConfig dense_config = config;
+    dense_config.num_experts = 0;
+    dense_config.top_k = 0;
+    dense_config.moe_intermediate = 0;
+    dense_config.expert_count = 0;
+    dense_config.intermediate_size = intermediate;
+    ctx = qwen36_create_training_context(
+        weight_ptrs.data(), static_cast<int64_t>(weight_ptrs.size()),
+        &embed, &final_norm, &lm_head, &dense_config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank,
+        &target_layer, 1, "q_proj");
+    assert(ctx);
+    if (world > 1) assert(qwen36_init_nccl(ctx) == 0);
+    const char* dense_targets = "gate_proj,up_proj,down_proj";
+    constexpr double fast_tenant_lr = 1e-2;
+    constexpr double default_tenant_lr = 1e-3;
+    constexpr double restored_tenant_lr = 5e-4;
+    constexpr double slow_tenant_lr = 1e-4;
+    constexpr double fast_beta1 = 0.8;
+    constexpr double fast_beta2 = 0.95;
+    constexpr double fast_eps = 1e-6;
+    constexpr double restored_beta1 = 0.7;
+    constexpr double restored_beta2 = 0.9;
+    constexpr double restored_eps = 1e-5;
+    const int64_t dense_fast = qwen36_add_lora_with_optimizer_v2(
+        ctx, rank, 1.0, &target_layer, 1, dense_targets,
+        fast_tenant_lr, fast_beta1, fast_beta2, fast_eps);
+    const int64_t dense_default = qwen36_add_lora(
+        ctx, rank, 1.0, &target_layer, 1, dense_targets);
+    const int64_t dense_slow = qwen36_add_lora_with_optimizer(
+        ctx, rank, 1.0, &target_layer, 1, dense_targets, slow_tenant_lr);
+    const int64_t dense_restored =
+        qwen36_add_lora_for_restore_with_optimizer_v2(
+            ctx, rank, 1.0, &target_layer, 1, dense_targets,
+            restored_tenant_lr, restored_beta1, restored_beta2, restored_eps);
+    assert(dense_fast > 0 && dense_default > dense_fast &&
+        dense_slow > dense_default && dense_restored > dense_slow);
+    assert(qwen36_add_lora_with_optimizer(
+        ctx, rank, 1.0, &target_layer, 1, dense_targets, NAN) < 0);
+    assert(qwen36_add_lora_with_optimizer_v2(
+        ctx, rank, 1.0, &target_layer, 1, dense_targets,
+        1e-3, 1.0, 0.999, 1e-8) < 0);
+    assert(qwen36_add_lora_with_optimizer_v2(
+        ctx, rank, 1.0, &target_layer, 1, dense_targets,
+        1e-3, std::nextafter(1.0, 0.0), 0.999, 1e-8) < 0);
+
+    auto* dense_fast_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(ctx, dense_fast, 0, "gate_proj", 0));
+    auto* dense_fast_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(ctx, dense_fast, 0, "gate_proj", 1));
+    assert(dense_fast_a && dense_fast_b && dense_fast_b->sizes() ==
+        at::IntArrayRef({intermediate, local_lora_rank}));
+    auto dense_a_value = dense_fast_a->clone();
+    auto dense_b_value = at::full(
+        dense_fast_b->sizes(), 0.01, dense_fast_b->options());
+    for (const int64_t adapter_id :
+         {dense_fast, dense_default, dense_slow, dense_restored}) {
+        assert(qwen36_set_adapter_lora_tensor(
+            ctx, adapter_id, 0, "gate_proj", 0, &dense_a_value) == 0);
+        assert(qwen36_set_adapter_lora_tensor(
+            ctx, adapter_id, 0, "gate_proj", 1, &dense_b_value) == 0);
+    }
+    const auto dense_b_before = dense_b_value.clone();
+    const int64_t dense_ids[] = {
+        dense_fast, dense_default, dense_slow, dense_restored};
+    const int64_t dense_adam_launches_before =
+        qwen36_get_dynamic_adam_launch_count(ctx);
+    const double dense_loss = qwen36_train_multi_lora_selected_v2(
+        ctx, &shared_input, &shared_target_row, &shared_attention,
+        dense_ids, 4);
+    c10::cuda::device_synchronize();
+    assert(qwen36_get_dynamic_adam_launch_count(ctx) ==
+        dense_adam_launches_before + 1);
+
+    auto validate_tenant_adam = [&](int64_t adapter_id, double learning_rate,
+                                     double beta1, double beta2, double eps) {
+        auto* parameter = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_lora_tensor(
+                ctx, adapter_id, 0, "gate_proj", 1));
+        auto* moment = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_optimizer_tensor(
+                ctx, adapter_id, 0, "gate_proj", 1, 0));
+        auto* variance = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_optimizer_tensor(
+                ctx, adapter_id, 0, "gate_proj", 1, 1));
+        assert(parameter && moment && variance);
+        const float beta1_f = static_cast<float>(beta1);
+        const float beta2_f = static_cast<float>(beta2);
+        const float step_lr = static_cast<float>(
+            learning_rate * std::sqrt(1.0 - std::pow(beta2_f, 1.0)) /
+            (1.0 - std::pow(beta1_f, 1.0)));
+        const float step_eps = static_cast<float>(
+            eps * std::sqrt(1.0 - std::pow(beta2_f, 1.0)));
+        auto expected = (dense_b_before.to(at::kFloat) - step_lr *
+            moment->to(at::kFloat) /
+            (variance->to(at::kFloat).sqrt() + step_eps)).to(at::kBFloat16);
+        assert((*parameter - expected).abs().max().item<double>() == 0.0);
+        return (*parameter - dense_b_before).abs().sum().item<double>();
+    };
+    const double fast_update = validate_tenant_adam(
+        dense_fast, fast_tenant_lr, fast_beta1, fast_beta2, fast_eps);
+    const double default_update =
+        validate_tenant_adam(dense_default, default_tenant_lr, 0.9, 0.999, 1e-8);
+    const double restored_update =
+        validate_tenant_adam(dense_restored, restored_tenant_lr,
+            restored_beta1, restored_beta2, restored_eps);
+    const double slow_update =
+        validate_tenant_adam(dense_slow, slow_tenant_lr, 0.9, 0.999, 1e-8);
+    auto normalized_tenant_state = [&](int64_t adapter_id,
+                                       double beta1, double beta2) {
+        auto* moment = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_optimizer_tensor(
+                ctx, adapter_id, 0, "gate_proj", 1, 0));
+        auto* variance = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_optimizer_tensor(
+                ctx, adapter_id, 0, "gate_proj", 1, 1));
+        assert(moment && variance);
+        return std::make_pair(
+            moment->to(at::kFloat) / (1.0 - beta1),
+            variance->to(at::kFloat) / (1.0 - beta2));
+    };
+    const auto default_state = normalized_tenant_state(
+        dense_default, 0.9, 0.999);
+    for (const auto& [adapter_id, beta1, beta2] :
+         std::vector<std::tuple<int64_t, double, double>>{
+             {dense_fast, fast_beta1, fast_beta2},
+             {dense_slow, 0.9, 0.999},
+             {dense_restored, restored_beta1, restored_beta2}}) {
+        const auto state = normalized_tenant_state(adapter_id, beta1, beta2);
+        assert(at::allclose(state.first, default_state.first, 2e-5, 1e-7));
+        assert(at::allclose(state.second, default_state.second, 2e-5, 1e-7));
+    }
+    std::printf(
+        "native_qwen35_tenant_optimizer_smoke loss=%0.8f "
+        "fast=%0.8e default=%0.8e restored=%0.8e slow=%0.8e\n",
+        dense_loss, fast_update, default_update, restored_update, slow_update);
+    assert(dense_loss == dense_loss && dense_loss > 0.0);
+    assert(fast_update > default_update && default_update > restored_update &&
+        restored_update > slow_update);
+    qwen36_free_training_context(ctx);
+
+    // Qwen3.5/3.6 linear-attention layers use the model's actual 128-wide
+    // delta-rule state. Exercise both fixed and per-sample dynamic GDN LoRA,
+    // including the custom CUDA backward for q/k/v/g/beta.
+    constexpr int64_t linear_heads = 1;
+    constexpr int64_t linear_dim = 128;
+    constexpr int64_t linear_qkv = 3 * linear_dim;
+    std::vector<at::Tensor> linear_weights;
+    linear_weights.push_back(cuda_rand({hidden}));
+    linear_weights.push_back(cuda_rand({hidden}));
+    linear_weights.push_back(cuda_rand({linear_qkv, hidden}));
+    linear_weights.push_back(cuda_rand({linear_dim, hidden}));
+    linear_weights.push_back(cuda_rand({linear_heads, hidden}));
+    linear_weights.push_back(cuda_rand({linear_heads, hidden}));
+    linear_weights.push_back(cuda_rand({linear_heads}));
+    linear_weights.push_back(cuda_rand({linear_heads}));
+    linear_weights.push_back(cuda_rand({linear_qkv, 1, 4}));
+    linear_weights.push_back(cuda_rand({linear_dim}));
+    linear_weights.push_back(cuda_rand({hidden, linear_dim}));
+    linear_weights.push_back(cuda_rand({intermediate, hidden}));
+    linear_weights.push_back(cuda_rand({intermediate, hidden}));
+    linear_weights.push_back(cuda_rand({hidden, intermediate}));
+    for (auto& weight : linear_weights) weight.set_requires_grad(false);
+    weight_ptrs.clear();
+    for (auto& weight : linear_weights) weight_ptrs.push_back(&weight);
+    LayerConfig linear_config = dense_config;
+    linear_config.layer_type = 1;
+    linear_config.num_heads = 0;
+    linear_config.num_kv_heads = 0;
+    linear_config.head_dim = 0;
+    linear_config.num_k_heads = linear_heads;
+    linear_config.key_dim = linear_dim;
+    linear_config.num_v_heads = linear_heads;
+    linear_config.val_dim = linear_dim;
+    linear_config.conv_kernel = 4;
+    ctx = qwen36_create_training_context(
+        weight_ptrs.data(), static_cast<int64_t>(weight_ptrs.size()),
+        &embed, &final_norm, &lm_head, &linear_config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank,
+        &target_layer, 1,
+        "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj");
+    assert(ctx);
+    if (world > 1) assert(qwen36_init_nccl(ctx) == 0);
+    assert(qwen36_get_lora_count(ctx) == 8);
+    // A chunk smaller than the causal-convolution overlap used to form a
+    // negative narrow() start for the second chunk. Reject it before any
+    // collective or tensor allocation instead of failing deep in ATen.
+    auto invalid_chunk_input = at::arange(
+        1, 5, at::TensorOptions().device(at::kCUDA).dtype(at::kLong))
+        .reshape({1, 4});
+    auto invalid_chunk_target = at::ones(
+        {1, 4}, at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    auto invalid_chunk_attention = at::ones(
+        {1, 4}, at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+    setenv("QWEN36_SEQ_CHUNK", "2", 1);
+    const double invalid_chunk_loss = qwen36_eval_step(
+        ctx, &invalid_chunk_input, &invalid_chunk_target,
+        &invalid_chunk_attention);
+    unsetenv("QWEN36_SEQ_CHUNK");
+    assert(invalid_chunk_loss < 0.0);
+    auto* linear_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(ctx, 0));
+    auto* linear_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(ctx, 0));
+    assert(linear_a && linear_b);
+    assert(linear_a->sizes() == at::IntArrayRef({local_lora_rank, hidden}));
+    assert(linear_b->sizes() == at::IntArrayRef({linear_qkv, local_lora_rank}));
+    auto linear_b_value = at::ones(linear_b->sizes(), linear_b->options());
+    assert(qwen36_set_lora_tensor(ctx, 0, 1, &linear_b_value) == 0);
+    auto linear_a_before = linear_a->clone();
+    auto* linear_a_accum = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(ctx, 0, 0));
+    assert(linear_a_accum);
+    assert(linear_a_accum->scalar_type() == at::kFloat);
+    assert(linear_a_accum->sizes() == linear_a->sizes());
+    assert(linear_a_accum->abs().sum().item<double>() == 0.0);
+    assert(qwen36_get_step_count(ctx) == 0);
+    assert(qwen36_set_step_count(ctx, -1) != 0);
+    assert(qwen36_get_step_count(ctx) == 0);
+    assert(qwen36_set_step_count(ctx, 7) == 0);
+    assert(qwen36_get_step_count(ctx) == 7);
+    assert(qwen36_set_step_count(ctx, 0) == 0);
+    const double accum_loss_0 = qwen36_train_micro_step(
+        ctx, &input_ids, &target_mask, &attention_mask, 0.5, 0);
+    c10::cuda::device_synchronize();
+    assert(accum_loss_0 == accum_loss_0);
+    assert(qwen36_get_step_count(ctx) == 0);
+    assert((*linear_a - linear_a_before).abs().sum().item<double>() == 0.0);
+    assert(linear_a_accum->abs().sum().item<double>() > 0.0);
+    assert(qwen36_get_accumulation_active(ctx) == 1);
+    const double pending_token_weight =
+        qwen36_get_accumulated_token_weight(ctx);
+    assert(pending_token_weight > 0.0);
+    assert(qwen36_add_lora(
+        ctx, rank, 1.0, &target_layer, 1, "in_proj_qkv") < 0);
+    assert(qwen36_add_lora_v2(
+        ctx, rank, 1.0, &target_layer, 1, "in_proj_qkv") < 0);
+    assert(qwen36_add_lora_for_restore(
+        ctx, rank, 1.0, &target_layer, 1, "in_proj_qkv") < 0);
+    assert(qwen36_set_adapter_id(ctx, 1, 2) < 0);
+    assert(qwen36_remove_lora(ctx, 1) < 0);
+    // The BF16 leaf gradient is consumed at the micro-step boundary; only the
+    // FP32 accumulator may survive.
+    assert(!linear_a->grad().defined());
+    auto pending_accumulator = linear_a_accum->clone();
+    const int64_t unavailable_dynamic_id[] = {1};
+    assert(qwen36_train_multi_lora_selected_v2(
+        ctx, &input_ids, &target_mask, &attention_mask,
+        unavailable_dynamic_id, 1) < 0.0);
+    c10::cuda::device_synchronize();
+    assert((*linear_a_accum - pending_accumulator).abs()
+        .max().item<double>() == 0.0);
+    assert(qwen36_get_accumulation_active(ctx) == 1);
+    assert(qwen36_get_accumulated_token_weight(ctx) ==
+        pending_token_weight);
+
+    // A failing final micro-step aborts the existing window transactionally.
+    const double failed_accum = qwen36_train_micro_step(
+        ctx, &input_ids, &target_mask, &attention_mask, NAN, 1);
+    c10::cuda::device_synchronize();
+    assert(failed_accum < 0.0);
+    assert(qwen36_get_step_count(ctx) == 0);
+    assert(linear_a_accum->abs().sum().item<double>() == 0.0);
+    assert((*linear_a - linear_a_before).abs().sum().item<double>() == 0.0);
+
+    // Explicit abort is idempotent and clears a successful non-final micro.
+    assert(qwen36_train_micro_step(
+        ctx, &input_ids, &target_mask, &attention_mask, 0.5, 0) > 0.0);
+    c10::cuda::device_synchronize();
+    assert(linear_a_accum->abs().sum().item<double>() > 0.0);
+    assert(qwen36_abort_gradient_accumulation(ctx) == 0);
+    c10::cuda::device_synchronize();
+    assert(linear_a_accum->abs().sum().item<double>() == 0.0);
+    assert(qwen36_get_step_count(ctx) == 0);
+    assert((*linear_a - linear_a_before).abs().sum().item<double>() == 0.0);
+
+    // A globally empty fixed window is a strict no-op: no Adam state/clock
+    // transition and no lingering accumulation state.
+    auto fixed_zero_target_mask = at::zeros_like(target_mask);
+    const auto linear_a_before_empty = linear_a->clone();
+    const double empty_window_loss = qwen36_train_micro_step(
+        ctx, &input_ids, &fixed_zero_target_mask, &attention_mask, 0.5, 1);
+    c10::cuda::device_synchronize();
+    assert(std::isfinite(empty_window_loss));
+    assert(qwen36_get_step_count(ctx) == 0);
+    assert((*linear_a - linear_a_before_empty).abs()
+        .max().item<double>() == 0.0);
+    assert(linear_a_accum->abs().sum().item<double>() == 0.0);
+    assert(qwen36_get_accumulation_active(ctx) == 0);
+    assert(qwen36_get_accumulated_token_weight(ctx) == 0.0);
+
+    // Positive, zero, positive micro-batches accumulate into FP32 and commit
+    // one Adam step. The zero-token micro must not dilute the numerator or
+    // denominator already present in the window.
+    const double clean_accum_loss_0 = qwen36_train_micro_step(
+        ctx, &input_ids, &target_mask, &attention_mask, 0.5, 0);
+    c10::cuda::device_synchronize();
+    assert(clean_accum_loss_0 == clean_accum_loss_0);
+    assert(linear_a_accum->scalar_type() == at::kFloat);
+    assert(linear_a_accum->abs().sum().item<double>() > 0.0);
+    const auto positive_accumulator = linear_a_accum->clone();
+    const double positive_token_weight =
+        qwen36_get_accumulated_token_weight(ctx);
+    const double zero_micro_loss = qwen36_train_micro_step(
+        ctx, &input_ids, &fixed_zero_target_mask, &attention_mask, 0.5, 0);
+    c10::cuda::device_synchronize();
+    assert(std::isfinite(zero_micro_loss));
+    assert(qwen36_get_step_count(ctx) == 0);
+    assert((*linear_a_accum - positive_accumulator).abs()
+        .max().item<double>() == 0.0);
+    assert(qwen36_get_accumulated_token_weight(ctx) ==
+        positive_token_weight);
+    const double accum_loss_1 = qwen36_train_micro_step(
+        ctx, &input_ids, &target_mask, &attention_mask, 0.5, 1);
+    c10::cuda::device_synchronize();
+    assert(accum_loss_1 == accum_loss_1);
+    assert(qwen36_get_step_count(ctx) == 1);
+    assert((*linear_a - linear_a_before).abs().sum().item<double>() > 0.0);
+    assert(linear_a_accum->abs().sum().item<double>() == 0.0);
+
+    // A non-finite accumulated gradient must abort before the in-place Adam
+    // launch. Parameters, optimizer state, and the bias-correction clock are
+    // one transaction; none may advance on rejection.
+    const auto poisoned_a_before = linear_a->clone();
+    const auto poisoned_b_before = linear_b->clone();
+    const int64_t optimizer_count = qwen36_get_lora_count(ctx) * 2;
+    std::vector<void*> optimizer_m_ptrs(optimizer_count);
+    std::vector<void*> optimizer_v_ptrs(optimizer_count);
+    assert(qwen36_export_optimizer_state(
+        ctx, optimizer_m_ptrs.data(), optimizer_v_ptrs.data(),
+        optimizer_count) == optimizer_count);
+    std::vector<at::Tensor> optimizer_m_before;
+    std::vector<at::Tensor> optimizer_v_before;
+    optimizer_m_before.reserve(optimizer_count);
+    optimizer_v_before.reserve(optimizer_count);
+    for (int64_t index = 0; index < optimizer_count; ++index) {
+        optimizer_m_before.push_back(
+            reinterpret_cast<at::Tensor*>(optimizer_m_ptrs[index])->clone());
+        optimizer_v_before.push_back(
+            reinterpret_cast<at::Tensor*>(optimizer_v_ptrs[index])->clone());
+    }
+    linear_a_accum->fill_(NAN);
+    const double poisoned_loss = qwen36_train_micro_step(
+        ctx, &input_ids, &target_mask, &attention_mask, 1.0, 1);
+    c10::cuda::device_synchronize();
+    assert(poisoned_loss < 0.0);
+    assert(qwen36_get_step_count(ctx) == 1);
+    assert(at::equal(*linear_a, poisoned_a_before));
+    assert(at::equal(*linear_b, poisoned_b_before));
+    assert(linear_a_accum->abs().sum().item<double>() == 0.0);
+    for (int64_t index = 0; index < optimizer_count; ++index) {
+        assert(at::equal(
+            *reinterpret_cast<at::Tensor*>(optimizer_m_ptrs[index]),
+            optimizer_m_before[index]));
+        assert(at::equal(
+            *reinterpret_cast<at::Tensor*>(optimizer_v_ptrs[index]),
+            optimizer_v_before[index]));
+    }
+
+    assert(qwen36_set_max_grad_norm(ctx, 1.0e-4) == 0);
+    linear_a_before = linear_a->clone();
+    const double linear_loss = qwen36_train_step(
+        ctx, &input_ids, &target_mask, &attention_mask);
+    c10::cuda::device_synchronize();
+    const double linear_update =
+        (*linear_a - linear_a_before).abs().sum().item<double>();
+    std::printf(
+        "native_qwen35_linear_lora_smoke loss=%0.8f qkv_a_update=%0.8e\n",
+        linear_loss, linear_update);
+    assert(linear_loss == linear_loss && linear_loss > 0.0);
+    assert(linear_update > 0.0);
+    assert(qwen36_set_max_grad_norm(ctx, 0.0) == 0);
+
+    const int64_t linear_adapter_one = qwen36_add_lora(
+        ctx, rank, 1.0, &target_layer, 1, "in_proj_qkv");
+    const int64_t linear_adapter_two = qwen36_add_lora(
+        ctx, rank, 1.0, &target_layer, 1, "in_proj_qkv");
+    assert(linear_adapter_one > 0 && linear_adapter_two > linear_adapter_one);
+    auto* dynamic_linear_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, linear_adapter_one, 0, "in_proj_qkv", 1));
+    assert(dynamic_linear_b && dynamic_linear_b->sizes() ==
+        at::IntArrayRef({linear_qkv, local_lora_rank}));
+    auto dynamic_linear_b_value = at::full(
+        dynamic_linear_b->sizes(), 0.01, dynamic_linear_b->options());
+    assert(qwen36_set_adapter_lora_tensor(
+        ctx, linear_adapter_one, 0, "in_proj_qkv", 1,
+        &dynamic_linear_b_value) == 0);
+    auto dynamic_linear_b_before = dynamic_linear_b->clone();
+    assert(qwen36_set_max_grad_norm(ctx, 1.0e-4) == 0);
+    const double dynamic_linear_loss = qwen36_train_multi_lora(
+        ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask, 2, rank);
+    c10::cuda::device_synchronize();
+    dynamic_linear_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            ctx, linear_adapter_one, 0, "in_proj_qkv", 1));
+    assert(dynamic_linear_b);
+    const double dynamic_linear_update =
+        (*dynamic_linear_b - dynamic_linear_b_before).abs().sum().item<double>();
+    std::printf(
+        "native_qwen35_linear_multi_lora_smoke loss=%0.8f qkv_b_update=%0.8e\n",
+        dynamic_linear_loss, dynamic_linear_update);
+    assert(dynamic_linear_loss == dynamic_linear_loss && dynamic_linear_loss > 0.0);
+    assert(qwen36_set_max_grad_norm(ctx, 0.0) == 0);
+    assert(dynamic_linear_update > 0.0);
+    qwen36_free_training_context(ctx);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_tp_attention_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_attention_smoke.cpp
new file mode 100644
index 00000000..31024967
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_tp_attention_smoke.cpp
@@ -0,0 +1,563 @@
+#include <ATen/ATen.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_nccl(void*);
+extern "C" int32_t qwen36_set_mtp_weights(
+    void*, void*, void*, void*, void*, void**, int64_t, void*, int64_t);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t);
+extern "C" int32_t qwen36_abort_gradient_accumulation(void*);
+extern "C" int64_t qwen36_export_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" double qwen36_eval_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_micro_step(
+    void*, void*, void*, void*, double, int32_t);
+extern "C" double qwen36_train_multi_lora_selected(
+    void*, void*, void*, void*, const int64_t*, int32_t, int32_t);
+extern "C" void qwen36_free_training_context(void*);
+
+static at::Tensor deterministic(std::initializer_list<int64_t> shape, double scale) {
+    int64_t count = 1;
+    for (int64_t dim : shape) count *= dim;
+    return ((at::arange(count,
+                 at::TensorOptions().device(at::kCUDA).dtype(at::kFloat))
+                 .remainder(19) - 9.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) {
+    return (lhs - rhs).abs().max().item<double>();
+}
+
+int main() {
+    const int rank = std::atoi(std::getenv("RANK"));
+    const int world = std::atoi(std::getenv("WORLD_SIZE"));
+    assert(world == 2 && rank >= 0 && rank < world);
+    qwen36_set_cuda_device(rank);
+    auto set_distributed_env = [&]() {
+        setenv("WORLD_SIZE", std::to_string(world).c_str(), 1);
+        setenv("RANK", std::to_string(rank).c_str(), 1);
+        setenv("TP_SIZE", "2", 1);
+        setenv("CP_SIZE", "1", 1);
+        setenv("EP_SIZE", "1", 1);
+        setenv("DP_SIZE", "1", 1);
+        setenv("PP_SIZE", "1", 1);
+        setenv("RUSTRAIN_TP_RANK", std::to_string(rank).c_str(), 1);
+    };
+    auto set_reference_env = []() {
+        setenv("WORLD_SIZE", "1", 1);
+        setenv("RANK", "0", 1);
+        setenv("TP_SIZE", "1", 1);
+        setenv("CP_SIZE", "1", 1);
+        setenv("EP_SIZE", "1", 1);
+        setenv("DP_SIZE", "1", 1);
+        setenv("PP_SIZE", "1", 1);
+        setenv("RUSTRAIN_TP_RANK", "0", 1);
+    };
+
+    constexpr int64_t hidden = 8;
+    constexpr int64_t heads = 4;
+    constexpr int64_t kv_heads = 2;
+    constexpr int64_t head_dim = 2;
+    constexpr int64_t intermediate = 12;
+    constexpr int64_t vocab = 16;
+    constexpr int64_t lora_rank = 4;
+
+    std::vector<at::Tensor> full_weights;
+    full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({2 * heads * head_dim, hidden}, 0.010));
+    full_weights.push_back(at::ones({head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({kv_heads * head_dim, hidden}, 0.012));
+    full_weights.push_back(at::ones({head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({kv_heads * head_dim, hidden}, 0.008));
+    full_weights.push_back(deterministic({hidden, heads * head_dim}, 0.011));
+    full_weights.push_back(deterministic({intermediate, hidden}, 0.009));
+    full_weights.push_back(deterministic({intermediate, hidden}, 0.007));
+    full_weights.push_back(deterministic({hidden, intermediate}, 0.010));
+    for (auto& weight : full_weights) weight.set_requires_grad(false);
+
+    const int64_t local_heads = heads / world;
+    const int64_t local_kv_heads = kv_heads / world;
+    std::vector<at::Tensor> local_weights;
+    local_weights.push_back(full_weights[0]);
+    local_weights.push_back(full_weights[1]);
+    local_weights.push_back(full_weights[2].narrow(
+        0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim).contiguous());
+    local_weights.push_back(full_weights[3]);
+    local_weights.push_back(full_weights[4].narrow(
+        0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim).contiguous());
+    local_weights.push_back(full_weights[5]);
+    local_weights.push_back(full_weights[6].narrow(
+        0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim).contiguous());
+    local_weights.push_back(full_weights[7].narrow(
+        1, rank * local_heads * head_dim, local_heads * head_dim).contiguous());
+    local_weights.insert(
+        local_weights.end(), full_weights.begin() + 8, full_weights.end());
+
+    auto embed = deterministic({vocab, hidden}, 0.020);
+    auto final_norm = at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+    auto lm_head = deterministic({vocab, hidden}, 0.015);
+    LayerConfig config{};
+    config.layer_type = 0;
+    config.num_heads = heads;
+    config.num_kv_heads = kv_heads;
+    config.head_dim = head_dim;
+    config.partial_rotary_factor = 1.0;
+    config.rope_theta = 10000.0;
+    config.rms_eps = 1e-5;
+    config.intermediate_size = intermediate;
+
+    const int64_t target_layer = 0;
+    auto local_ptrs = pointers(local_weights);
+    set_distributed_env();
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    void* distributed = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, "q_proj,k_proj,v_proj,o_proj", 1);
+    assert(distributed);
+    assert(qwen36_set_mtp_weights(
+        distributed, nullptr, nullptr, nullptr, nullptr, nullptr, 0, nullptr, 0) != 0);
+    assert(qwen36_init_nccl(distributed) == 0);
+
+    set_reference_env();
+    auto full_ptrs = pointers(full_weights);
+    void* reference = qwen36_create_training_context(
+        full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, "q_proj,k_proj,v_proj,o_proj");
+    assert(reference);
+    set_distributed_env();
+
+    auto q_a = deterministic({lora_rank, hidden}, 0.0020);
+    auto q_b = deterministic({2 * heads * head_dim, lora_rank}, 0.0010);
+    auto k_a = deterministic({lora_rank, hidden}, 0.0018);
+    auto k_b = deterministic({kv_heads * head_dim, lora_rank}, 0.0011);
+    auto v_a = deterministic({lora_rank, hidden}, 0.0016);
+    auto v_b = deterministic({kv_heads * head_dim, lora_rank}, 0.0009);
+    auto o_a = deterministic({lora_rank, heads * head_dim}, 0.0015);
+    auto o_b = deterministic({hidden, lora_rank}, 0.0012);
+    auto local_q_b = q_b.narrow(
+        0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim).contiguous();
+    auto local_o_a = o_a.narrow(
+        1, rank * local_heads * head_dim, local_heads * head_dim).contiguous();
+    auto local_k_b = k_b.narrow(
+        0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim).contiguous();
+    auto local_v_b = v_b.narrow(
+        0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim).contiguous();
+    assert(qwen36_set_lora_tensor(distributed, 0, 0, &q_a) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 0, 1, &local_q_b) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 1, 0, &k_a) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 1, 1, &local_k_b) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 2, 0, &v_a) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 2, 1, &local_v_b) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 3, 0, &local_o_a) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 3, 1, &o_b) == 0);
+    assert(qwen36_set_lora_tensor(reference, 0, 0, &q_a) == 0);
+    assert(qwen36_set_lora_tensor(reference, 0, 1, &q_b) == 0);
+    assert(qwen36_set_lora_tensor(reference, 1, 0, &k_a) == 0);
+    assert(qwen36_set_lora_tensor(reference, 1, 1, &k_b) == 0);
+    assert(qwen36_set_lora_tensor(reference, 2, 0, &v_a) == 0);
+    assert(qwen36_set_lora_tensor(reference, 2, 1, &v_b) == 0);
+    assert(qwen36_set_lora_tensor(reference, 3, 0, &o_a) == 0);
+    assert(qwen36_set_lora_tensor(reference, 3, 1, &o_b) == 0);
+
+    auto input_ids = at::tensor({1, 2, 3},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 3});
+    auto target_mask = at::ones({1, 3},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    auto attention_mask = at::ones({1, 3},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+
+    const double distributed_eval = qwen36_eval_step(
+        distributed, &input_ids, &target_mask, &attention_mask);
+    const double reference_eval = qwen36_eval_step(
+        reference, &input_ids, &target_mask, &attention_mask);
+    assert(distributed_eval > 0.0 && reference_eval > 0.0);
+
+    const double distributed_micro = qwen36_train_micro_step(
+        distributed, &input_ids, &target_mask, &attention_mask, 1.0, 0);
+    const double reference_micro = qwen36_train_micro_step(
+        reference, &input_ids, &target_mask, &attention_mask, 1.0, 0);
+    assert(distributed_micro > 0.0 && reference_micro > 0.0);
+    auto* local_q_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(distributed, 0, 1));
+    auto* full_q_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(reference, 0, 1));
+    auto* local_k_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(distributed, 1, 1));
+    auto* full_k_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(reference, 1, 1));
+    auto* local_v_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(distributed, 2, 1));
+    auto* full_v_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(reference, 2, 1));
+    auto* local_o_a_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(distributed, 3, 0));
+    auto* full_o_a_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(reference, 3, 0));
+    assert(local_q_b_grad && full_q_b_grad && local_k_b_grad && full_k_b_grad);
+    assert(local_v_b_grad && full_v_b_grad && local_o_a_grad && full_o_a_grad);
+    const double q_b_grad_diff = max_diff(
+        *local_q_b_grad,
+        full_q_b_grad->narrow(
+            0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim));
+    const double o_a_grad_diff = max_diff(
+        *local_o_a_grad,
+        full_o_a_grad->narrow(
+            1, rank * local_heads * head_dim, local_heads * head_dim));
+    const double k_b_grad_diff = max_diff(
+        *local_k_b_grad,
+        full_k_b_grad->narrow(
+            0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim));
+    const double v_b_grad_diff = max_diff(
+        *local_v_b_grad,
+        full_v_b_grad->narrow(
+            0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim));
+    assert(qwen36_abort_gradient_accumulation(distributed) == 0);
+    assert(qwen36_abort_gradient_accumulation(reference) == 0);
+
+    const double distributed_loss = qwen36_train_step(
+        distributed, &input_ids, &target_mask, &attention_mask);
+    const double reference_loss = qwen36_train_step(
+        reference, &input_ids, &target_mask, &attention_mask);
+    assert(distributed_loss > 0.0 && reference_loss > 0.0);
+
+    auto* updated_q_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(distributed, 0));
+    auto* updated_q_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(distributed, 0));
+    auto* updated_k_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(distributed, 1));
+    auto* updated_k_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(distributed, 1));
+    auto* updated_v_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(distributed, 2));
+    auto* updated_v_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(distributed, 2));
+    auto* updated_o_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(distributed, 3));
+    auto* updated_o_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(distributed, 3));
+    auto* reference_q_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(reference, 0));
+    auto* reference_q_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(reference, 0));
+    auto* reference_k_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(reference, 1));
+    auto* reference_k_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(reference, 1));
+    auto* reference_v_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(reference, 2));
+    auto* reference_v_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(reference, 2));
+    auto* reference_o_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(reference, 3));
+    auto* reference_o_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(reference, 3));
+    assert(updated_q_a && updated_q_b && updated_k_a && updated_k_b);
+    assert(updated_v_a && updated_v_b && updated_o_a && updated_o_b);
+    assert(reference_q_a && reference_q_b && reference_k_a && reference_k_b);
+    assert(reference_v_a && reference_v_b && reference_o_a && reference_o_b);
+
+    constexpr int64_t optimizer_count = 14;
+    std::vector<void*> local_m(optimizer_count), local_v(optimizer_count);
+    std::vector<void*> full_m(optimizer_count), full_v(optimizer_count);
+    assert(qwen36_export_optimizer_state(
+        distributed, local_m.data(), local_v.data(), optimizer_count) ==
+        optimizer_count);
+    assert(qwen36_export_optimizer_state(
+        reference, full_m.data(), full_v.data(), optimizer_count) ==
+        optimizer_count);
+    double optimizer_m_diff = 0.0;
+    double optimizer_v_diff = 0.0;
+    double adam_error = 0.0;
+    auto observe_optimizer = [&](int64_t index, const at::Tensor& expected_m,
+                                 const at::Tensor& expected_v,
+                                 const at::Tensor& updated,
+                                 const at::Tensor& before) {
+        auto* m = reinterpret_cast<at::Tensor*>(local_m[index]);
+        auto* v = reinterpret_cast<at::Tensor*>(local_v[index]);
+        assert(m && v);
+        optimizer_m_diff = std::max(optimizer_m_diff, max_diff(*m, expected_m));
+        optimizer_v_diff = std::max(optimizer_v_diff, max_diff(*v, expected_v));
+        auto expected_param = (
+            before.to(at::kFloat) - 1e-3 *
+                (*m / (1.0 - 0.9)) /
+                (((*v / (1.0 - 0.999)).sqrt()) + 1e-8))
+            .to(at::kBFloat16);
+        adam_error = std::max(adam_error, max_diff(updated, expected_param));
+    };
+    auto state = [](std::vector<void*>& tensors, int64_t index) -> at::Tensor& {
+        auto* tensor = reinterpret_cast<at::Tensor*>(tensors[index]);
+        assert(tensor);
+        return *tensor;
+    };
+    observe_optimizer(0, state(full_m, 0), state(full_v, 0), *updated_q_a, q_a);
+    observe_optimizer(1,
+        state(full_m, 1).narrow(
+            0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim),
+        state(full_v, 1).narrow(
+            0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim),
+        *updated_q_b, local_q_b);
+    observe_optimizer(2, state(full_m, 2), state(full_v, 2), *updated_k_a, k_a);
+    observe_optimizer(3,
+        state(full_m, 3).narrow(
+            0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim),
+        state(full_v, 3).narrow(
+            0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim),
+        *updated_k_b, local_k_b);
+    observe_optimizer(4, state(full_m, 4), state(full_v, 4), *updated_v_a, v_a);
+    observe_optimizer(5,
+        state(full_m, 5).narrow(
+            0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim),
+        state(full_v, 5).narrow(
+            0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim),
+        *updated_v_b, local_v_b);
+    observe_optimizer(6,
+        state(full_m, 6).narrow(
+            1, rank * local_heads * head_dim, local_heads * head_dim),
+        state(full_v, 6).narrow(
+            1, rank * local_heads * head_dim, local_heads * head_dim),
+        *updated_o_a, local_o_a);
+    observe_optimizer(7, state(full_m, 7), state(full_v, 7), *updated_o_b, o_b);
+    const double q_a_diff = max_diff(*updated_q_a, *reference_q_a);
+    const double q_b_diff = max_diff(
+        *updated_q_b,
+        reference_q_b->narrow(
+            0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim));
+    const double o_a_diff = max_diff(
+        *updated_o_a,
+        reference_o_a->narrow(
+            1, rank * local_heads * head_dim, local_heads * head_dim));
+    const double o_b_diff = max_diff(*updated_o_b, *reference_o_b);
+    const double k_a_diff = max_diff(*updated_k_a, *reference_k_a);
+    const double k_b_diff = max_diff(
+        *updated_k_b, reference_k_b->narrow(
+            0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim));
+    const double v_a_diff = max_diff(*updated_v_a, *reference_v_a);
+    const double v_b_diff = max_diff(
+        *updated_v_b, reference_v_b->narrow(
+            0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim));
+    std::printf(
+        "base_tp_attention_smoke rank=%d eval_diff=%0.8e loss_diff=%0.8e q_b_grad_diff=%0.8e k_b_grad_diff=%0.8e v_b_grad_diff=%0.8e o_a_grad_diff=%0.8e m_diff=%0.8e v_diff=%0.8e adam_error=%0.8e q_a_diff=%0.8e q_b_diff=%0.8e k_a_diff=%0.8e k_b_diff=%0.8e v_a_diff=%0.8e v_b_diff=%0.8e o_a_diff=%0.8e o_b_diff=%0.8e\n",
+        rank, std::abs(distributed_eval - reference_eval),
+        std::abs(distributed_loss - reference_loss), q_b_grad_diff,
+        k_b_grad_diff, v_b_grad_diff, o_a_grad_diff, optimizer_m_diff,
+        optimizer_v_diff, adam_error, q_a_diff, q_b_diff, k_a_diff,
+        k_b_diff, v_a_diff, v_b_diff, o_a_diff, o_b_diff);
+    std::fflush(stdout);
+    // Row-parallel BF16 rounds each local matmul before the NCCL sum, while
+    // the reference rounds once after the full matmul.
+    assert(std::abs(distributed_eval - reference_eval) < 5e-3);
+    assert(std::abs(distributed_loss - reference_loss) < 5e-3);
+    assert(q_b_grad_diff < 3e-4 && k_b_grad_diff < 3e-4);
+    assert(v_b_grad_diff < 5e-4 && o_a_grad_diff < 5e-4);
+    assert(optimizer_m_diff < 5e-5 && optimizer_v_diff < 5e-8);
+    assert(adam_error < 1e-8);
+    assert(std::max({q_a_diff, q_b_diff, k_a_diff, k_b_diff,
+        v_a_diff, v_b_diff, o_a_diff, o_b_diff}) <= 2e-3);
+
+    const char* fused_lora_qkv = std::getenv("QWEN36_FUSED_LORA_QKV_A");
+    if (fused_lora_qkv && std::string(fused_lora_qkv) != "0") {
+        setenv("QWEN36_FUSED_LORA_QKV_A", "0", 1);
+        const double legacy_lora_qkv = qwen36_eval_step(
+            distributed, &input_ids, &target_mask, &attention_mask);
+        setenv("QWEN36_FUSED_LORA_QKV_A", "1", 1);
+        const double fused_lora_qkv_loss = qwen36_eval_step(
+            distributed, &input_ids, &target_mask, &attention_mask);
+        assert(std::abs(fused_lora_qkv_loss - legacy_lora_qkv) < 5e-5);
+
+        if (rank == 1) setenv("QWEN36_FUSED_LORA_QKV_A", "0", 1);
+        const double flag_mismatch = qwen36_eval_step(
+            distributed, &input_ids, &target_mask, &attention_mask);
+        assert(flag_mismatch < 0.0);
+        setenv("QWEN36_FUSED_LORA_QKV_A", "1", 1);
+    }
+
+    qwen36_free_training_context(reference);
+    qwen36_free_training_context(distributed);
+
+    set_distributed_env();
+    void* dynamic_distributed = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, "v_proj", 1);
+    assert(dynamic_distributed && qwen36_init_nccl(dynamic_distributed) == 0);
+    const int64_t distributed_id = qwen36_add_lora(
+        dynamic_distributed, lora_rank, lora_rank,
+        &target_layer, 1, "q_proj,k_proj,v_proj,o_proj");
+    assert(distributed_id > 0);
+
+    set_reference_env();
+    void* dynamic_reference = qwen36_create_training_context(
+        full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, "v_proj");
+    assert(dynamic_reference);
+    const int64_t reference_id = qwen36_add_lora(
+        dynamic_reference, lora_rank, lora_rank,
+        &target_layer, 1, "q_proj,k_proj,v_proj,o_proj");
+    assert(reference_id > 0);
+    set_distributed_env();
+
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_distributed, distributed_id, 0, "q_proj", 0, &q_a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_distributed, distributed_id, 0, "q_proj", 1, &local_q_b) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_distributed, distributed_id, 0, "k_proj", 0, &k_a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_distributed, distributed_id, 0, "k_proj", 1, &local_k_b) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_distributed, distributed_id, 0, "v_proj", 0, &v_a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_distributed, distributed_id, 0, "v_proj", 1, &local_v_b) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_distributed, distributed_id, 0, "o_proj", 0, &local_o_a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_distributed, distributed_id, 0, "o_proj", 1, &o_b) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_reference, reference_id, 0, "q_proj", 0, &q_a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_reference, reference_id, 0, "q_proj", 1, &q_b) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_reference, reference_id, 0, "k_proj", 0, &k_a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_reference, reference_id, 0, "k_proj", 1, &k_b) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_reference, reference_id, 0, "v_proj", 0, &v_a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_reference, reference_id, 0, "v_proj", 1, &v_b) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_reference, reference_id, 0, "o_proj", 0, &o_a) == 0);
+    assert(qwen36_set_adapter_lora_tensor(
+        dynamic_reference, reference_id, 0, "o_proj", 1, &o_b) == 0);
+
+    const double dynamic_distributed_loss = qwen36_train_multi_lora_selected(
+        dynamic_distributed, &input_ids, &target_mask, &attention_mask,
+        &distributed_id, 1, lora_rank);
+    const double dynamic_reference_loss = qwen36_train_multi_lora_selected(
+        dynamic_reference, &input_ids, &target_mask, &attention_mask,
+        &reference_id, 1, lora_rank);
+    assert(dynamic_distributed_loss > 0.0 && dynamic_reference_loss > 0.0);
+
+    auto* dynamic_q_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_distributed, distributed_id, 0, "q_proj", 0));
+    auto* dynamic_q_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_distributed, distributed_id, 0, "q_proj", 1));
+    auto* dynamic_k_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_distributed, distributed_id, 0, "k_proj", 0));
+    auto* dynamic_k_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_distributed, distributed_id, 0, "k_proj", 1));
+    auto* dynamic_v_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_distributed, distributed_id, 0, "v_proj", 0));
+    auto* dynamic_v_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_distributed, distributed_id, 0, "v_proj", 1));
+    auto* dynamic_o_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_distributed, distributed_id, 0, "o_proj", 0));
+    auto* dynamic_o_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_distributed, distributed_id, 0, "o_proj", 1));
+    auto* dynamic_reference_q_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_reference, reference_id, 0, "q_proj", 0));
+    auto* dynamic_reference_q_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_reference, reference_id, 0, "q_proj", 1));
+    auto* dynamic_reference_k_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_reference, reference_id, 0, "k_proj", 0));
+    auto* dynamic_reference_k_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_reference, reference_id, 0, "k_proj", 1));
+    auto* dynamic_reference_v_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_reference, reference_id, 0, "v_proj", 0));
+    auto* dynamic_reference_v_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_reference, reference_id, 0, "v_proj", 1));
+    auto* dynamic_reference_o_a = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_reference, reference_id, 0, "o_proj", 0));
+    auto* dynamic_reference_o_b = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            dynamic_reference, reference_id, 0, "o_proj", 1));
+    assert(dynamic_q_a && dynamic_q_b && dynamic_o_a && dynamic_o_b);
+    assert(dynamic_k_a && dynamic_k_b && dynamic_v_a && dynamic_v_b);
+    assert(dynamic_reference_q_a && dynamic_reference_q_b &&
+        dynamic_reference_o_a && dynamic_reference_o_b);
+    assert(dynamic_reference_k_a && dynamic_reference_k_b &&
+        dynamic_reference_v_a && dynamic_reference_v_b);
+    const double dynamic_q_a_diff = max_diff(*dynamic_q_a, *dynamic_reference_q_a);
+    const double dynamic_q_b_diff = max_diff(
+        *dynamic_q_b, dynamic_reference_q_b->narrow(
+            0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim));
+    const double dynamic_o_a_diff = max_diff(
+        *dynamic_o_a, dynamic_reference_o_a->narrow(
+            1, rank * local_heads * head_dim, local_heads * head_dim));
+    const double dynamic_o_b_diff = max_diff(*dynamic_o_b, *dynamic_reference_o_b);
+    const double dynamic_k_a_diff = max_diff(*dynamic_k_a, *dynamic_reference_k_a);
+    const double dynamic_k_b_diff = max_diff(
+        *dynamic_k_b, dynamic_reference_k_b->narrow(
+            0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim));
+    const double dynamic_v_a_diff = max_diff(*dynamic_v_a, *dynamic_reference_v_a);
+    const double dynamic_v_b_diff = max_diff(
+        *dynamic_v_b, dynamic_reference_v_b->narrow(
+            0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim));
+    std::printf(
+        "dynamic_tp_attention_smoke rank=%d loss_diff=%0.8e q_a_diff=%0.8e q_b_diff=%0.8e k_a_diff=%0.8e k_b_diff=%0.8e v_a_diff=%0.8e v_b_diff=%0.8e o_a_diff=%0.8e o_b_diff=%0.8e\n",
+        rank,
+        std::abs(dynamic_distributed_loss - dynamic_reference_loss),
+        dynamic_q_a_diff, dynamic_q_b_diff, dynamic_k_a_diff, dynamic_k_b_diff,
+        dynamic_v_a_diff, dynamic_v_b_diff, dynamic_o_a_diff, dynamic_o_b_diff);
+    std::fflush(stdout);
+    assert(std::abs(dynamic_distributed_loss - dynamic_reference_loss) < 5e-3);
+    assert(dynamic_q_a_diff < 5e-5 && dynamic_q_b_diff < 5e-5);
+    assert(dynamic_k_a_diff < 5e-5 && dynamic_k_b_diff < 5e-5);
+    assert(dynamic_v_a_diff < 5e-5 && dynamic_v_b_diff < 5e-5);
+    assert(dynamic_o_a_diff < 5e-5 && dynamic_o_b_diff < 5e-5);
+
+    qwen36_free_training_context(dynamic_reference);
+    qwen36_free_training_context(dynamic_distributed);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_tp_dp_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_dp_smoke.cpp
new file mode 100644
index 00000000..194f74fd
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_tp_dp_smoke.cpp
@@ -0,0 +1,555 @@
+#include <ATen/ATen.h>
+
+#include <algorithm>
+#include <array>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" int64_t qwen36_kernel_abi_version();
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_nccl(void*);
+extern "C" int32_t qwen36_init_parallel_nccl(
+    void*, int32_t, int32_t, int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" int64_t qwen36_export_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" double qwen36_train_multi_lora_selected(
+    void*, void*, void*, void*, const int64_t*, int32_t, int32_t);
+extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t);
+extern "C" void qwen36_free_training_context(void*);
+
+static constexpr double kLearningRate = 1e-3;
+static constexpr double kBeta1 = 0.9;
+static constexpr double kBeta2 = 0.999;
+static constexpr double kAdamEps = 1e-8;
+static constexpr int64_t kOptimizerSlots = 14;
+static constexpr int32_t kBaseTpAttention = 1 << 0;
+static constexpr int32_t kDataParallel = 1 << 1;
+static constexpr int32_t kVocabParallel = 1 << 2;
+
+static at::Tensor deterministic(std::initializer_list<int64_t> shape, double scale) {
+    int64_t count = 1;
+    for (int64_t dim : shape) count *= dim;
+    return ((at::arange(count,
+                 at::TensorOptions().device(at::kCUDA).dtype(at::kFloat))
+                 .remainder(19) - 9.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) {
+    return (lhs - rhs).abs().max().item<double>();
+}
+
+struct Batch {
+    at::Tensor input_ids;
+    at::Tensor target_mask;
+    at::Tensor attention_mask;
+};
+
+static Batch make_batch(int dp_rank) {
+    auto options_long = at::TensorOptions().device(at::kCUDA).dtype(at::kLong);
+    auto options_float = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat);
+    auto options_bool = at::TensorOptions().device(at::kCUDA).dtype(at::kBool);
+    if (dp_rank == 0) {
+        return {
+            at::tensor({1, 2, 3, 4}, options_long).reshape({1, 4}),
+            at::tensor({1.0, 1.0, 1.0, 1.0}, options_float).reshape({1, 4}),
+            at::ones({1, 4}, options_bool),
+        };
+    }
+    return {
+        at::tensor({4, 2, 5, 1}, options_long).reshape({1, 4}),
+        at::tensor({1.0, 1.0, 0.0, 0.0}, options_float).reshape({1, 4}),
+        at::ones({1, 4}, options_bool),
+    };
+}
+
+struct LoraWeights {
+    at::Tensor q_a, q_b;
+    at::Tensor k_a, k_b;
+    at::Tensor v_a, v_b;
+    at::Tensor o_a, o_b;
+};
+
+static LoraWeights make_lora_weights(
+    int64_t hidden, int64_t heads, int64_t kv_heads,
+    int64_t head_dim, int64_t lora_rank
+) {
+    return {
+        deterministic({lora_rank, hidden}, 0.0020),
+        deterministic({2 * heads * head_dim, lora_rank}, 0.0010),
+        deterministic({lora_rank, hidden}, 0.0018),
+        deterministic({kv_heads * head_dim, lora_rank}, 0.0011),
+        deterministic({lora_rank, hidden}, 0.0016),
+        deterministic({kv_heads * head_dim, lora_rank}, 0.0009),
+        deterministic({lora_rank, heads * head_dim}, 0.0015),
+        deterministic({hidden, lora_rank}, 0.0012),
+    };
+}
+
+static LoraWeights clone_lora_weights(const LoraWeights& weights) {
+    return {
+        weights.q_a.clone(), weights.q_b.clone(),
+        weights.k_a.clone(), weights.k_b.clone(),
+        weights.v_a.clone(), weights.v_b.clone(),
+        weights.o_a.clone(), weights.o_b.clone(),
+    };
+}
+
+static void offset_lora_weights(LoraWeights& weights, double offset) {
+    std::array<at::Tensor*, 8> tensors = {
+        &weights.q_a, &weights.q_b, &weights.k_a, &weights.k_b,
+        &weights.v_a, &weights.v_b, &weights.o_a, &weights.o_b,
+    };
+    for (auto* tensor : tensors) tensor->add_(offset);
+}
+
+static void install_full_lora(void* context, LoraWeights& weights) {
+    std::array<at::Tensor*, 8> tensors = {
+        &weights.q_a, &weights.q_b, &weights.k_a, &weights.k_b,
+        &weights.v_a, &weights.v_b, &weights.o_a, &weights.o_b,
+    };
+    for (int64_t index = 0; index < 4; ++index) {
+        assert(qwen36_set_lora_tensor(context, index, 0, tensors[2 * index]) == 0);
+        assert(qwen36_set_lora_tensor(context, index, 1, tensors[2 * index + 1]) == 0);
+    }
+}
+
+static LoraWeights shard_lora_for_tp(
+    const LoraWeights& full, int tp_rank, int64_t local_heads,
+    int64_t local_kv_heads, int64_t head_dim
+) {
+    return {
+        full.q_a,
+        full.q_b.narrow(
+            0, tp_rank * 2 * local_heads * head_dim,
+            2 * local_heads * head_dim).contiguous(),
+        full.k_a,
+        full.k_b.narrow(
+            0, tp_rank * local_kv_heads * head_dim,
+            local_kv_heads * head_dim).contiguous(),
+        full.v_a,
+        full.v_b.narrow(
+            0, tp_rank * local_kv_heads * head_dim,
+            local_kv_heads * head_dim).contiguous(),
+        full.o_a.narrow(
+            1, tp_rank * local_heads * head_dim,
+            local_heads * head_dim).contiguous(),
+        full.o_b,
+    };
+}
+
+static std::array<at::Tensor*, 8> lora_parameters(void* context) {
+    std::array<at::Tensor*, 8> result{};
+    for (int64_t index = 0; index < 4; ++index) {
+        result[2 * index] = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(context, index));
+        result[2 * index + 1] = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(context, index));
+        assert(result[2 * index] && result[2 * index + 1]);
+    }
+    return result;
+}
+
+struct OptimizerState {
+    std::vector<void*> m;
+    std::vector<void*> v;
+};
+
+static OptimizerState optimizer_state(void* context) {
+    OptimizerState result{
+        std::vector<void*>(kOptimizerSlots),
+        std::vector<void*>(kOptimizerSlots),
+    };
+    assert(qwen36_export_optimizer_state(
+        context, result.m.data(), result.v.data(), kOptimizerSlots) ==
+        kOptimizerSlots);
+    return result;
+}
+
+static at::Tensor& state_tensor(std::vector<void*>& tensors, int64_t index) {
+    auto* tensor = reinterpret_cast<at::Tensor*>(tensors[index]);
+    assert(tensor);
+    return *tensor;
+}
+
+struct FullReference {
+    void* context;
+    double loss;
+    OptimizerState optimizer;
+};
+
+static FullReference run_full_reference(
+    std::vector<void*>& full_ptrs, at::Tensor& embed, at::Tensor& final_norm,
+    at::Tensor& lm_head, LayerConfig& config, LoraWeights& lora,
+    Batch& batch, int64_t vocab, int64_t lora_rank
+) {
+    const int64_t target_layer = 0;
+    void* context = qwen36_create_training_context(
+        full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        vocab, 1e-5, lora_rank, &target_layer, 1,
+        "q_proj,k_proj,v_proj,o_proj");
+    assert(context);
+    install_full_lora(context, lora);
+    const double loss = qwen36_train_step(
+        context, &batch.input_ids, &batch.target_mask, &batch.attention_mask);
+    assert(loss > 0.0 && std::isfinite(loss));
+    return {context, loss, optimizer_state(context)};
+}
+
+int main() {
+    assert(qwen36_kernel_abi_version() == 31);
+    const int rank = std::atoi(std::getenv("RANK"));
+    const int world = std::atoi(std::getenv("WORLD_SIZE"));
+    const int local_rank = std::atoi(
+        std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0");
+    assert(world == 4 && rank >= 0 && rank < world);
+    // Deliberately use dp-tp rank order so vocabulary ownership follows the
+    // explicit topology rank rather than global_rank % TP_SIZE.
+    const int tp_rank = rank / 2;
+    const int dp_rank = rank % 2;
+    qwen36_set_cuda_device(local_rank);
+
+    constexpr int64_t hidden = 8;
+    constexpr int64_t heads = 4;
+    constexpr int64_t kv_heads = 2;
+    constexpr int64_t head_dim = 2;
+    constexpr int64_t intermediate = 12;
+    constexpr int64_t vocab = 16;
+    constexpr int64_t lora_rank = 4;
+    constexpr int64_t local_heads = heads / 2;
+    constexpr int64_t local_kv_heads = kv_heads / 2;
+
+    std::vector<at::Tensor> full_weights;
+    full_weights.push_back(at::ones(
+        {hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(at::ones(
+        {hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({2 * heads * head_dim, hidden}, 0.010));
+    full_weights.push_back(at::ones(
+        {head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({kv_heads * head_dim, hidden}, 0.012));
+    full_weights.push_back(at::ones(
+        {head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({kv_heads * head_dim, hidden}, 0.008));
+    full_weights.push_back(deterministic({hidden, heads * head_dim}, 0.011));
+    full_weights.push_back(deterministic({intermediate, hidden}, 0.009));
+    full_weights.push_back(deterministic({intermediate, hidden}, 0.007));
+    full_weights.push_back(deterministic({hidden, intermediate}, 0.010));
+    for (auto& weight : full_weights) weight.set_requires_grad(false);
+
+    std::vector<at::Tensor> local_weights;
+    local_weights.push_back(full_weights[0]);
+    local_weights.push_back(full_weights[1]);
+    local_weights.push_back(full_weights[2].narrow(
+        0, tp_rank * 2 * local_heads * head_dim,
+        2 * local_heads * head_dim).contiguous());
+    local_weights.push_back(full_weights[3]);
+    local_weights.push_back(full_weights[4].narrow(
+        0, tp_rank * local_kv_heads * head_dim,
+        local_kv_heads * head_dim).contiguous());
+    local_weights.push_back(full_weights[5]);
+    local_weights.push_back(full_weights[6].narrow(
+        0, tp_rank * local_kv_heads * head_dim,
+        local_kv_heads * head_dim).contiguous());
+    local_weights.push_back(full_weights[7].narrow(
+        1, tp_rank * local_heads * head_dim,
+        local_heads * head_dim).contiguous());
+    local_weights.insert(
+        local_weights.end(), full_weights.begin() + 8, full_weights.end());
+
+    auto embed = deterministic({vocab, hidden}, 0.020);
+    auto final_norm = at::ones(
+        {hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+    auto lm_head = deterministic({vocab, hidden}, 0.015);
+    auto local_embed = embed.narrow(
+        0, tp_rank * (vocab / 2), vocab / 2).contiguous();
+    auto local_lm_head = lm_head.narrow(
+        0, tp_rank * (vocab / 2), vocab / 2).contiguous();
+    LayerConfig config{};
+    config.layer_type = 0;
+    config.num_heads = heads;
+    config.num_kv_heads = kv_heads;
+    config.head_dim = head_dim;
+    config.partial_rotary_factor = 1.0;
+    config.rope_theta = 10000.0;
+    config.rms_eps = 1e-5;
+    config.intermediate_size = intermediate;
+
+    auto local_ptrs = pointers(local_weights);
+    auto full_ptrs = pointers(full_weights);
+    const int64_t target_layer = 0;
+
+    // Native validation must reject a topology whose declared axes do not
+    // cover the process world before any communicator is initialized.
+    setenv("TP_SIZE", "2", 1);
+    setenv("DP_SIZE", "3", 1);
+    setenv("RUSTRAIN_DATA_PARALLEL", "1", 1);
+    void* invalid = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        vocab, 1e-5, lora_rank, &target_layer, 1,
+        "q_proj,k_proj,v_proj,o_proj",
+        kBaseTpAttention | kDataParallel | kVocabParallel);
+    assert(invalid == nullptr);
+
+    setenv("DP_SIZE", "2", 1);
+    auto full_lora = make_lora_weights(
+        hidden, heads, kv_heads, head_dim, lora_rank);
+    auto local_lora = shard_lora_for_tp(
+        full_lora, tp_rank, local_heads, local_kv_heads, head_dim);
+    auto synchronized_lora = clone_lora_weights(local_lora);
+    if (dp_rank == 1) {
+        // DP must restore every local shard from the matching dp_rank=0 peer.
+        offset_lora_weights(local_lora, 0.25);
+    } else if (tp_rank == 1) {
+        // TP must restore projection-replicated factors from tp_rank=0.
+        local_lora.q_a.add_(0.25);
+        local_lora.k_a.add_(0.25);
+        local_lora.v_a.add_(0.25);
+        local_lora.o_b.add_(0.25);
+    }
+    void* distributed = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        vocab, 1e-5, lora_rank, &target_layer, 1,
+        "q_proj,k_proj,v_proj,o_proj",
+        kBaseTpAttention | kDataParallel | kVocabParallel);
+    assert(distributed);
+    install_full_lora(distributed, local_lora);
+    assert(qwen36_init_parallel_nccl(
+        distributed, rank, world,
+        tp_rank, 2, dp_rank,
+        0, 1, rank,
+        dp_rank, 2, tp_rank) == 0);
+
+    const auto synchronized_params = lora_parameters(distributed);
+    const std::array<at::Tensor*, 8> expected_synchronized = {
+        &synchronized_lora.q_a, &synchronized_lora.q_b,
+        &synchronized_lora.k_a, &synchronized_lora.k_b,
+        &synchronized_lora.v_a, &synchronized_lora.v_b,
+        &synchronized_lora.o_a, &synchronized_lora.o_b,
+    };
+    double broadcast_diff = 0.0;
+    for (int64_t index = 0; index < 8; ++index) {
+        broadcast_diff = std::max(
+            broadcast_diff,
+            max_diff(*synchronized_params[index], *expected_synchronized[index]));
+    }
+    assert(broadcast_diff == 0.0);
+
+    auto batch0 = make_batch(0);
+    auto batch1 = make_batch(1);
+    Batch full_batch{
+        at::cat({batch0.input_ids, batch1.input_ids}, 0),
+        at::cat({batch0.target_mask, batch1.target_mask}, 0),
+        at::cat({batch0.attention_mask, batch1.attention_mask}, 0),
+    };
+    Batch& local_batch = dp_rank == 0 ? batch0 : batch1;
+    const double distributed_loss = qwen36_train_step(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask);
+    assert(distributed_loss > 0.0 && std::isfinite(distributed_loss));
+
+    // Reference contexts stay process-local. The concatenated batch is the
+    // hard oracle for the DP-reduced update; the two local references prove
+    // that unequal token counts are not accidentally averaged by replica.
+    setenv("WORLD_SIZE", "1", 1);
+    setenv("RANK", "0", 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    auto reference_lora = make_lora_weights(
+        hidden, heads, kv_heads, head_dim, lora_rank);
+    auto reference = run_full_reference(
+        full_ptrs, embed, final_norm, lm_head, config, reference_lora,
+        full_batch, vocab, lora_rank);
+    auto local0_lora = make_lora_weights(
+        hidden, heads, kv_heads, head_dim, lora_rank);
+    auto local0 = run_full_reference(
+        full_ptrs, embed, final_norm, lm_head, config, local0_lora,
+        batch0, vocab, lora_rank);
+    auto local1_lora = make_lora_weights(
+        hidden, heads, kv_heads, head_dim, lora_rank);
+    auto local1 = run_full_reference(
+        full_ptrs, embed, final_norm, lm_head, config, local1_lora,
+        batch1, vocab, lora_rank);
+
+    auto distributed_params = lora_parameters(distributed);
+    auto reference_params = lora_parameters(reference.context);
+    auto distributed_optimizer = optimizer_state(distributed);
+    std::array<at::Tensor, 8> before = {
+        synchronized_lora.q_a, synchronized_lora.q_b,
+        synchronized_lora.k_a, synchronized_lora.k_b,
+        synchronized_lora.v_a, synchronized_lora.v_b,
+        synchronized_lora.o_a, synchronized_lora.o_b,
+    };
+    auto reference_slice = [&](const at::Tensor& tensor, int64_t index) {
+        switch (index) {
+            case 1:
+                return tensor.narrow(
+                    0, tp_rank * 2 * local_heads * head_dim,
+                    2 * local_heads * head_dim);
+            case 3:
+            case 5:
+                return tensor.narrow(
+                    0, tp_rank * local_kv_heads * head_dim,
+                    local_kv_heads * head_dim);
+            case 6:
+                return tensor.narrow(
+                    1, tp_rank * local_heads * head_dim,
+                    local_heads * head_dim);
+            default:
+                return tensor;
+        }
+    };
+
+    double parameter_diff = 0.0;
+    double optimizer_m_diff = 0.0;
+    double optimizer_v_diff = 0.0;
+    double adam_error = 0.0;
+    double weighted_reference_error = 0.0;
+    double equal_replica_gap = 0.0;
+    constexpr double token_count0 = 3.0;
+    constexpr double token_count1 = 1.0;
+    for (int64_t index = 0; index < 8; ++index) {
+        auto expected_parameter = reference_slice(*reference_params[index], index);
+        auto expected_m = reference_slice(
+            state_tensor(reference.optimizer.m, index), index);
+        auto expected_v = reference_slice(
+            state_tensor(reference.optimizer.v, index), index);
+        parameter_diff = std::max(
+            parameter_diff, max_diff(*distributed_params[index], expected_parameter));
+        optimizer_m_diff = std::max(
+            optimizer_m_diff,
+            max_diff(state_tensor(distributed_optimizer.m, index), expected_m));
+        optimizer_v_diff = std::max(
+            optimizer_v_diff,
+            max_diff(state_tensor(distributed_optimizer.v, index), expected_v));
+
+        auto expected_adam_parameter = (
+            before[index].to(at::kFloat) - kLearningRate *
+                (state_tensor(distributed_optimizer.m, index) / (1.0 - kBeta1)) /
+                ((state_tensor(distributed_optimizer.v, index) /
+                    (1.0 - kBeta2)).sqrt() + kAdamEps))
+            .to(at::kBFloat16);
+        adam_error = std::max(
+            adam_error,
+            max_diff(*distributed_params[index], expected_adam_parameter));
+
+        auto local_m0 = state_tensor(local0.optimizer.m, index);
+        auto local_m1 = state_tensor(local1.optimizer.m, index);
+        auto token_weighted_m =
+            (local_m0 * token_count0 + local_m1 * token_count1) /
+            (token_count0 + token_count1);
+        auto equal_replica_m = (local_m0 + local_m1) * 0.5;
+        weighted_reference_error = std::max(
+            weighted_reference_error, max_diff(
+                state_tensor(reference.optimizer.m, index), token_weighted_m));
+        equal_replica_gap = std::max(
+            equal_replica_gap, max_diff(
+                state_tensor(reference.optimizer.m, index), equal_replica_m));
+    }
+
+    const double expected_local_loss =
+        dp_rank == 0 ? local0.loss : local1.loss;
+    const double local_loss_diff = std::abs(distributed_loss - expected_local_loss);
+    std::printf(
+        "native_tp_dp_smoke rank=%d tp_rank=%d dp_rank=%d "
+        "broadcast_diff=%0.8e local_loss_diff=%0.8e "
+        "parameter_diff=%0.8e m_diff=%0.8e "
+        "v_diff=%0.8e adam_error=%0.8e weighted_reference_error=%0.8e "
+        "equal_replica_gap=%0.8e\n",
+        rank, tp_rank, dp_rank, broadcast_diff, local_loss_diff, parameter_diff,
+        optimizer_m_diff, optimizer_v_diff, adam_error,
+        weighted_reference_error, equal_replica_gap);
+    std::fflush(stdout);
+
+    // TP row-parallel BF16 rounds each local matmul before the collective.
+    assert(local_loss_diff < 5e-3);
+    assert(parameter_diff <= 2e-3);
+    assert(optimizer_m_diff < 5e-5 && optimizer_v_diff < 5e-8);
+    assert(adam_error < 1e-8);
+    assert(weighted_reference_error < 5e-5);
+    assert(equal_replica_gap > 1e-8);
+
+    // A fixed-size registry signature must reject a different tenant order
+    // before any adapter-shaped collective can mix gradients or deadlock.
+    const int64_t adapter_one = qwen36_add_lora(
+        distributed, lora_rank, lora_rank, &target_layer, 1, "q_proj");
+    const int64_t adapter_two = qwen36_add_lora(
+        distributed, lora_rank, lora_rank, &target_layer, 1, "q_proj");
+    assert(adapter_one > 0 && adapter_two > adapter_one);
+    std::array<int64_t, 2> mismatched_ids = dp_rank == 0
+        ? std::array<int64_t, 2>{adapter_one, adapter_two}
+        : std::array<int64_t, 2>{adapter_two, adapter_one};
+    auto dynamic_input = local_batch.input_ids.repeat({2, 1});
+    auto dynamic_target = local_batch.target_mask.repeat({2, 1});
+    auto dynamic_attention = local_batch.attention_mask.repeat({2, 1});
+    const double mismatched_loss = qwen36_train_multi_lora_selected(
+        distributed, &dynamic_input, &dynamic_target, &dynamic_attention,
+        mismatched_ids.data(), mismatched_ids.size(), lora_rank);
+    assert(mismatched_loss < 0.0);
+    assert(qwen36_get_adapter_step_count(distributed, adapter_one) == 0);
+    assert(qwen36_get_adapter_step_count(distributed, adapter_two) == 0);
+
+    const std::array<int64_t, 2> ordered_ids{adapter_one, adapter_two};
+    const double dynamic_loss = qwen36_train_multi_lora_selected(
+        distributed, &dynamic_input, &dynamic_target, &dynamic_attention,
+        ordered_ids.data(), ordered_ids.size(), lora_rank);
+    assert(dynamic_loss > 0.0 && std::isfinite(dynamic_loss));
+    assert(qwen36_get_adapter_step_count(distributed, adapter_one) == 1);
+    assert(qwen36_get_adapter_step_count(distributed, adapter_two) == 1);
+
+    qwen36_free_training_context(local1.context);
+    qwen36_free_training_context(local0.context);
+    qwen36_free_training_context(reference.context);
+    qwen36_free_training_context(distributed);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_tp_ep_bench.cpp b/crates/rustrain-qwen3-6/tests/native_tp_ep_bench.cpp
new file mode 100644
index 00000000..02cd23ca
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_tp_ep_bench.cpp
@@ -0,0 +1,585 @@
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDACachingAllocator.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+#include <cassert>
+#include <chrono>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <iomanip>
+#include <initializer_list>
+#include <numeric>
+#include <sstream>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" int64_t qwen36_kernel_abi_version();
+extern "C" double qwen36_parallel_max_double(void*, double);
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_parallel_nccl(
+    void*, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_step_host_i64(
+    void*, const int64_t*, const int64_t*, const int64_t*, int64_t, int64_t);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int64_t qwen36_add_lora_v2(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" double qwen36_train_multi_lora_selected_v2(
+    void*, void*, void*, void*, const int64_t*, int32_t);
+extern "C" double qwen36_train_multi_lora_host_i64(
+    void*, const int64_t*, const int64_t*, const int64_t*, int64_t, int64_t,
+    int32_t, int32_t, const int64_t*, int32_t);
+extern "C" void qwen36_free_training_context(void*);
+
+namespace {
+
+constexpr int64_t kAbiVersion = 31;
+constexpr int32_t kBaseTpAttention = 1 << 0;
+constexpr int32_t kVocabParallel = 1 << 2;
+constexpr int32_t kExpertParallel = 1 << 3;
+constexpr int32_t kBaseTpMlp = 1 << 4;
+constexpr int kTpSize = 2;
+constexpr int kEpSize = 2;
+constexpr int kDpSize = 1;
+constexpr int kWorldSize = kTpSize * kEpSize * kDpSize;
+
+int env_int(const char* name, int fallback) {
+    const char* value = std::getenv(name);
+    if (!value || value[0] == '\0') return fallback;
+    const int parsed = std::atoi(value);
+    assert(parsed > 0);
+    return parsed;
+}
+
+int required_env_nonnegative(const char* name) {
+    const char* value = std::getenv(name);
+    assert(value && value[0] != '\0');
+    const int parsed = std::atoi(value);
+    assert(parsed >= 0);
+    return parsed;
+}
+
+bool env_enabled(const char* name, bool fallback = false) {
+    const char* value = std::getenv(name);
+    if (!value || value[0] == '\0') return fallback;
+    return std::strcmp(value, "0") != 0 && std::strcmp(value, "false") != 0;
+}
+
+std::string env_string(const char* name, const char* fallback) {
+    const char* value = std::getenv(name);
+    return value && value[0] != '\0' ? value : fallback;
+}
+
+std::string json_escape(const std::string& input) {
+    std::string output;
+    output.reserve(input.size());
+    for (const char value : input) {
+        switch (value) {
+            case '\\': output += "\\\\"; break;
+            case '"': output += "\\\""; break;
+            case '\n': output += "\\n"; break;
+            case '\r': output += "\\r"; break;
+            case '\t': output += "\\t"; break;
+            default: output += value; break;
+        }
+    }
+    return output;
+}
+
+at::Tensor seeded_cpu_randn(
+    std::initializer_list<int64_t> shape, double scale, int64_t seed
+) {
+    at::manual_seed(seed);
+    return (at::randn(shape,
+        at::TensorOptions().device(at::kCPU).dtype(at::kFloat)) * scale)
+        .to(at::kBFloat16);
+}
+
+at::Tensor cuda_local(const at::Tensor& tensor) {
+    return tensor.contiguous().to(at::kCUDA);
+}
+
+at::Tensor unit(std::initializer_list<int64_t> shape) {
+    return at::ones(
+        shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+}
+
+std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+void append_layer_weights(
+    std::vector<at::Tensor>& weights, int layer, int tp_rank, int ep_rank,
+    int hidden, int heads, int kv_heads, int head_dim,
+    int experts, int intermediate, bool expert_tp
+) {
+    const int local_heads = heads / kTpSize;
+    const int local_kv_heads = kv_heads / kTpSize;
+    const int local_experts = experts / kEpSize;
+    const int local_intermediate = expert_tp
+        ? intermediate / kTpSize : intermediate;
+    const int64_t common_seed = 1000 + layer * 100;
+
+    weights.push_back(unit({hidden}));
+    weights.push_back(unit({hidden}));
+    weights.push_back(cuda_local(seeded_cpu_randn(
+        {2 * heads * head_dim, hidden}, 0.0020, common_seed + 2).narrow(
+            0, tp_rank * 2 * local_heads * head_dim,
+            2 * local_heads * head_dim)));
+    weights.push_back(unit({head_dim}));
+    weights.push_back(cuda_local(seeded_cpu_randn(
+        {kv_heads * head_dim, hidden}, 0.0020, common_seed + 4).narrow(
+            0, tp_rank * local_kv_heads * head_dim,
+            local_kv_heads * head_dim)));
+    weights.push_back(unit({head_dim}));
+    weights.push_back(cuda_local(seeded_cpu_randn(
+        {kv_heads * head_dim, hidden}, 0.0020, common_seed + 6).narrow(
+            0, tp_rank * local_kv_heads * head_dim,
+            local_kv_heads * head_dim)));
+    weights.push_back(cuda_local(seeded_cpu_randn(
+        {hidden, heads * head_dim}, 0.0020, common_seed + 7).narrow(
+            1, tp_rank * local_heads * head_dim,
+            local_heads * head_dim)));
+    weights.push_back(cuda_local(seeded_cpu_randn(
+        {experts, hidden}, 0.0020, common_seed + 8)));
+    weights.push_back(cuda_local(seeded_cpu_randn(
+        {1, hidden}, 0.0020, common_seed + 9)));
+
+    auto shared_gate = seeded_cpu_randn(
+        {intermediate, hidden}, 0.0020, common_seed + 10);
+    auto shared_up = seeded_cpu_randn(
+        {intermediate, hidden}, 0.0020, common_seed + 11);
+    auto shared_down = seeded_cpu_randn(
+        {hidden, intermediate}, 0.0020, common_seed + 12);
+    if (expert_tp) {
+        shared_gate = shared_gate.narrow(
+            0, tp_rank * local_intermediate, local_intermediate);
+        shared_up = shared_up.narrow(
+            0, tp_rank * local_intermediate, local_intermediate);
+        shared_down = shared_down.narrow(
+            1, tp_rank * local_intermediate, local_intermediate);
+    }
+    weights.push_back(cuda_local(shared_gate));
+    weights.push_back(cuda_local(shared_up));
+    weights.push_back(cuda_local(shared_down));
+
+    auto expert_gate_up = seeded_cpu_randn(
+        {experts, 2 * intermediate, hidden}, 0.0020, common_seed + 13)
+        .narrow(0, ep_rank * local_experts, local_experts);
+    auto expert_down = seeded_cpu_randn(
+        {experts, hidden, intermediate}, 0.0020, common_seed + 14)
+        .narrow(0, ep_rank * local_experts, local_experts);
+    if (expert_tp) {
+        expert_gate_up = at::cat({
+            expert_gate_up.narrow(
+                1, tp_rank * local_intermediate, local_intermediate),
+            expert_gate_up.narrow(
+                1, intermediate + tp_rank * local_intermediate,
+                local_intermediate),
+        }, 1);
+        expert_down = expert_down.narrow(
+            2, tp_rank * local_intermediate, local_intermediate);
+    }
+    weights.push_back(cuda_local(expert_gate_up));
+    weights.push_back(cuda_local(expert_down));
+
+    assert(weights[weights.size() - 5].sizes() ==
+        at::IntArrayRef({local_intermediate, hidden}));
+    assert(weights[weights.size() - 4].sizes() ==
+        at::IntArrayRef({local_intermediate, hidden}));
+    assert(weights[weights.size() - 3].sizes() ==
+        at::IntArrayRef({hidden, local_intermediate}));
+    assert(weights[weights.size() - 2].sizes() ==
+        at::IntArrayRef({local_experts, 2 * local_intermediate, hidden}));
+    assert(weights[weights.size() - 1].sizes() ==
+        at::IntArrayRef({local_experts, hidden, local_intermediate}));
+}
+
+double percentile(std::vector<double> values, double quantile) {
+    assert(!values.empty());
+    std::sort(values.begin(), values.end());
+    const double position = quantile * static_cast<double>(values.size() - 1);
+    const size_t lower = static_cast<size_t>(position);
+    const size_t upper = std::min(lower + 1, values.size() - 1);
+    const double fraction = position - static_cast<double>(lower);
+    return values[lower] * (1.0 - fraction) + values[upper] * fraction;
+}
+
+double gib(int64_t bytes) {
+    return static_cast<double>(std::max<int64_t>(bytes, 0)) /
+        (1024.0 * 1024.0 * 1024.0);
+}
+
+size_t used_since(size_t initial_free, size_t observed_free) {
+    return initial_free > observed_free ? initial_free - observed_free : 0;
+}
+
+void set_parallel_env(int rank, int tp_rank, int ep_rank) {
+    setenv("WORLD_SIZE", "4", 1);
+    setenv("TP_SIZE", "2", 1);
+    setenv("EP_SIZE", "2", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", std::to_string(tp_rank).c_str(), 1);
+    setenv("RUSTRAIN_EP_RANK", std::to_string(ep_rank).c_str(), 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    setenv("RANK", std::to_string(rank).c_str(), 1);
+}
+
+}  // namespace
+
+int main() {
+    const int rank = required_env_nonnegative("RANK");
+    const int world = env_int("WORLD_SIZE", -1);
+    const int local_rank = required_env_nonnegative("LOCAL_RANK");
+    assert(world == kWorldSize && rank >= 0 && rank < world);
+    const int tp_rank = rank % kTpSize;
+    const int ep_rank = (rank / kTpSize) % kEpSize;
+    const int tp_color = ep_rank;
+    const int ep_color = tp_rank;
+    const int dp_color = rank;
+    assert(qwen36_kernel_abi_version() == kAbiVersion);
+    assert(env_enabled("QWEN36_EP_A2A"));
+    assert(env_enabled("QWEN36_EP_A2A_SHARDED"));
+    qwen36_set_cuda_device(local_rank);
+    assert(cudaFree(nullptr) == cudaSuccess);
+    set_parallel_env(rank, tp_rank, ep_rank);
+
+    const std::string lora_mode = env_string("BENCH_LORA_MODE", "fixed");
+    assert(lora_mode == "fixed" || lora_mode == "dynamic" ||
+        lora_mode == "heterogeneous");
+    const bool dynamic_lora = lora_mode != "fixed";
+    const std::string expert_tp_mode = env_string(
+        "BENCH_EXPERT_TP_MODE", "etp");
+    assert(expert_tp_mode == "replicated" || expert_tp_mode == "etp");
+    const bool expert_tp = expert_tp_mode == "etp";
+    const std::string variant = env_string("BENCH_VARIANT", "baseline");
+    const std::string input_abi = env_string("BENCH_INPUT_ABI", "tensor");
+    assert(input_abi == "tensor" || input_abi == "host");
+    const bool packed_a2a = env_enabled("QWEN36_EP_A2A_PACKED", true);
+    const std::string targets = env_string(
+        "BENCH_TARGETS", "q_proj,experts_gate_up_proj,experts_down_proj");
+    const std::string heterogeneous_targets = env_string(
+        "BENCH_HETERO_TARGETS", "q_proj");
+    const int batch = env_int("BENCH_BATCH", 2);
+    const int seq = env_int("BENCH_SEQ", 128);
+    const int hidden = env_int("BENCH_HIDDEN", 1024);
+    const int head_dim = env_int("BENCH_HEAD_DIM", 128);
+    const int heads = env_int("BENCH_HEADS", hidden / head_dim);
+    const int kv_heads = env_int("BENCH_KV_HEADS", heads);
+    const int experts = env_int("BENCH_EXPERTS", 8);
+    const int intermediate = env_int("BENCH_INTERMEDIATE", 2048);
+    const int top_k = env_int("BENCH_TOP_K", 2);
+    const bool gpu_metadata = env_enabled(
+        "QWEN36_EP_A2A_GPU_METADATA",
+        kEpSize >= 4 && batch * seq * (packed_a2a ? top_k : 1) >= 512);
+    const int lora_rank = env_int("BENCH_LORA_RANK", 16);
+    const int heterogeneous_rank = env_int(
+        "BENCH_HETERO_LORA_RANK", std::max(1, lora_rank / 2));
+    const int layers = env_int("BENCH_LAYERS", 1);
+    const int vocab = env_int("BENCH_VOCAB", 8192);
+    const int warmup = env_int("BENCH_WARMUP", 3);
+    const int iters = env_int("BENCH_ITERS", 20);
+    const int tenants = env_int("BENCH_TENANTS", std::max(batch, 8));
+    const int active_tenants = env_int("BENCH_ACTIVE_TENANTS", batch);
+    const bool rotate_tenants = env_enabled("BENCH_ROTATE_TENANTS");
+
+    assert(seq >= 2 && hidden > 0 && hidden == heads * head_dim);
+    assert(heads % kTpSize == 0 && kv_heads % kTpSize == 0);
+    assert(heads % kv_heads == 0);
+    assert(experts % kEpSize == 0 && top_k <= experts);
+    assert(lora_rank % kTpSize == 0 && vocab % kTpSize == 0);
+    assert(!expert_tp || intermediate % kTpSize == 0);
+    assert(!dynamic_lora ||
+        (active_tenants == batch && tenants >= active_tenants));
+
+    c10::cuda::CUDACachingAllocator::resetPeakStats(local_rank);
+    size_t free_start = 0;
+    size_t total_bytes = 0;
+    assert(cudaMemGetInfo(&free_start, &total_bytes) == cudaSuccess);
+    size_t min_observed_free = free_start;
+
+    std::vector<at::Tensor> weights;
+    weights.reserve(static_cast<size_t>(layers) * 15);
+    for (int layer = 0; layer < layers; ++layer) {
+        append_layer_weights(weights, layer, tp_rank, ep_rank,
+            hidden, heads, kv_heads, head_dim, experts, intermediate,
+            expert_tp);
+    }
+    for (auto& weight : weights) weight.set_requires_grad(false);
+    auto weight_ptrs = pointers(weights);
+
+    const int local_vocab = vocab / kTpSize;
+    auto embed = cuda_local(seeded_cpu_randn(
+        {vocab, hidden}, 0.0020, 41).narrow(
+            0, tp_rank * local_vocab, local_vocab));
+    auto final_norm = unit({hidden});
+    auto lm_head = cuda_local(seeded_cpu_randn(
+        {vocab, hidden}, 0.0020, 47).narrow(
+            0, tp_rank * local_vocab, local_vocab));
+    embed.set_requires_grad(false);
+    final_norm.set_requires_grad(false);
+    lm_head.set_requires_grad(false);
+
+    std::vector<LayerConfig> configs(layers);
+    for (auto& config : configs) {
+        config.layer_type = 0;
+        config.num_heads = heads;
+        config.num_kv_heads = kv_heads;
+        config.head_dim = head_dim;
+        config.partial_rotary_factor = 1.0;
+        config.rope_theta = 10000.0;
+        config.rms_eps = 1e-5;
+        config.num_experts = experts;
+        config.top_k = top_k;
+        config.moe_intermediate = intermediate;
+        config.expert_start = ep_rank * (experts / kEpSize);
+        config.expert_count = experts / kEpSize;
+        config.norm_topk_prob = 1;
+    }
+    std::vector<int64_t> target_layers(layers);
+    std::iota(target_layers.begin(), target_layers.end(), 0);
+
+    void* context = qwen36_create_training_context_ex(
+        weight_ptrs.data(), static_cast<int64_t>(weight_ptrs.size()),
+        &embed, &final_norm, &lm_head, configs.data(), layers,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        target_layers.data(), layers, targets.c_str(),
+        kBaseTpAttention | kVocabParallel | kExpertParallel |
+            (expert_tp ? kBaseTpMlp : 0));
+    assert(context);
+    assert(qwen36_init_parallel_nccl(
+        context, rank, world,
+        tp_rank, kTpSize, tp_color,
+        ep_rank, kEpSize, ep_color,
+        0, kDpSize, dp_color) == 0);
+
+    std::vector<int64_t> adapter_ids;
+    if (dynamic_lora) {
+        adapter_ids.reserve(tenants);
+        for (int tenant = 0; tenant < tenants; ++tenant) {
+            const bool heterogeneous_tenant =
+                lora_mode == "heterogeneous" && tenant % 2 != 0;
+            const int tenant_rank = heterogeneous_tenant
+                ? heterogeneous_rank : lora_rank;
+            const std::string& tenant_targets = heterogeneous_tenant
+                ? heterogeneous_targets : targets;
+            const int64_t adapter_id = lora_mode == "heterogeneous"
+                ? qwen36_add_lora_v2(
+                    context, tenant_rank, static_cast<double>(tenant_rank),
+                    target_layers.data(), layers, tenant_targets.c_str())
+                : qwen36_add_lora(
+                    context, tenant_rank, static_cast<double>(tenant_rank),
+                    target_layers.data(), layers, tenant_targets.c_str());
+            assert(adapter_id > 0);
+            adapter_ids.push_back(adapter_id);
+        }
+    }
+
+    std::vector<int64_t> host_ids(static_cast<size_t>(batch) * seq);
+    std::vector<int64_t> host_targets(static_cast<size_t>(batch) * seq, 1);
+    std::vector<int64_t> host_attention(static_cast<size_t>(batch) * seq, 1);
+    for (int b = 0; b < batch; ++b) {
+        for (int s = 0; s < seq; ++s) {
+            const int64_t global_source =
+                static_cast<int64_t>(ep_rank) * batch + b;
+            host_ids[static_cast<size_t>(b) * seq + s] =
+                1 + (global_source * 131 + s * 17) % (vocab - 1);
+        }
+    }
+    auto input_ids = at::from_blob(host_ids.data(), {batch, seq},
+        at::TensorOptions().device(at::kCPU).dtype(at::kLong))
+        .clone().to(at::kCUDA);
+    auto target_mask = at::ones({batch, seq},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    auto attention_mask = at::ones({batch, seq},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+
+    size_t free_after_context = 0;
+    assert(cudaMemGetInfo(&free_after_context, &total_bytes) == cudaSuccess);
+    min_observed_free = std::min(min_observed_free, free_after_context);
+
+    std::vector<int64_t> selected(active_tenants);
+    auto select_tenants = [&](int step) {
+        if (!dynamic_lora) return;
+        const int start = rotate_tenants
+            ? (step * active_tenants) % tenants : 0;
+        for (int index = 0; index < active_tenants; ++index) {
+            selected[index] = adapter_ids[(start + index) % tenants];
+        }
+    };
+    auto train_once = [&](int step) {
+        if (input_abi == "host") {
+            select_tenants(step);
+            if (lora_mode == "fixed") {
+                return qwen36_train_step_host_i64(
+                    context, host_ids.data(), host_targets.data(),
+                    host_attention.data(), batch, seq);
+            }
+            return qwen36_train_multi_lora_host_i64(
+                context, host_ids.data(), host_targets.data(),
+                host_attention.data(), batch, seq, active_tenants, lora_rank,
+                selected.data(), active_tenants);
+        }
+        if (lora_mode == "fixed") {
+            return qwen36_train_step(
+                context, &input_ids, &target_mask, &attention_mask);
+        }
+        select_tenants(step);
+        return qwen36_train_multi_lora_selected_v2(
+            context, &input_ids, &target_mask, &attention_mask,
+            selected.data(), active_tenants);
+    };
+
+    double last_loss = 0.0;
+    for (int step = 0; step < warmup; ++step) {
+        last_loss = train_once(step);
+        assert(last_loss > 0.0 && std::isfinite(last_loss));
+    }
+    assert(cudaDeviceSynchronize() == cudaSuccess);
+    size_t free_after_warmup = 0;
+    assert(cudaMemGetInfo(&free_after_warmup, &total_bytes) == cudaSuccess);
+    min_observed_free = std::min(min_observed_free, free_after_warmup);
+
+    std::vector<double> times_ms;
+    times_ms.reserve(iters);
+    for (int iteration = 0; iteration < iters; ++iteration) {
+        assert(cudaDeviceSynchronize() == cudaSuccess);
+        const auto start = std::chrono::steady_clock::now();
+        last_loss = train_once(warmup + iteration);
+        assert(cudaDeviceSynchronize() == cudaSuccess);
+        const auto stop = std::chrono::steady_clock::now();
+        assert(last_loss > 0.0 && std::isfinite(last_loss));
+        const double local_ms = std::chrono::duration<double, std::milli>(
+            stop - start).count();
+        const double world_max_ms = qwen36_parallel_max_double(context, local_ms);
+        assert(world_max_ms > 0.0 && std::isfinite(world_max_ms));
+        times_ms.push_back(world_max_ms);
+        size_t free_now = 0;
+        assert(cudaMemGetInfo(&free_now, &total_bytes) == cudaSuccess);
+        min_observed_free = std::min(min_observed_free, free_now);
+    }
+
+    const double mean = std::accumulate(
+        times_ms.begin(), times_ms.end(), 0.0) / times_ms.size();
+    double variance = 0.0;
+    for (const double value : times_ms) {
+        variance += (value - mean) * (value - mean);
+    }
+    variance /= times_ms.size();
+    const double p50 = percentile(times_ms, 0.50);
+    const double p95 = percentile(times_ms, 0.95);
+    const double unique_tokens =
+        static_cast<double>(batch) * seq * kEpSize;
+    const double loss_tokens =
+        static_cast<double>(batch) * (seq - 1) * kEpSize;
+    const double routed_tokens = unique_tokens * top_k * layers;
+    const double unique_tokens_per_sec = unique_tokens / (p50 / 1000.0);
+    const double routed_tokens_per_sec = routed_tokens / (p50 / 1000.0);
+
+    const auto allocator_stats =
+        c10::cuda::CUDACachingAllocator::getDeviceStats(local_rank);
+    // Aggregate is the stable first entry of CUDACachingAllocator StatArray;
+    // the enum namespace differs across prebuilt PyTorch releases.
+    constexpr size_t aggregate = 0;
+    const auto& allocated = allocator_stats.allocated_bytes[aggregate];
+    const auto& reserved = allocator_stats.reserved_bytes[aggregate];
+    cudaDeviceProp properties{};
+    assert(cudaGetDeviceProperties(&properties, local_rank) == cudaSuccess);
+
+    std::ostringstream output;
+    output << std::fixed << std::setprecision(6)
+        << "native_tp_ep_bench {"
+        << "\"variant\":\"" << json_escape(variant) << "\","
+        << "\"lora_mode\":\"" << lora_mode << "\","
+        << "\"input_abi\":\"" << input_abi << "\","
+        << "\"expert_tp_mode\":\"" << expert_tp_mode << "\","
+        << "\"targets\":\"" << json_escape(targets) << "\","
+        << "\"rank\":" << rank << ",\"world\":" << world << ','
+        << "\"tp_rank\":" << tp_rank << ",\"tp_size\":" << kTpSize << ','
+        << "\"ep_rank\":" << ep_rank << ",\"ep_size\":" << kEpSize << ','
+        << "\"gpu\":\"" << json_escape(properties.name) << "\","
+        << "\"abi\":" << kAbiVersion << ','
+        << "\"ep_a2a\":true,\"ep_a2a_sharded\":true,"
+        << "\"ep_a2a_packed\":" << (packed_a2a ? "true" : "false") << ','
+        << "\"ep_a2a_gpu_metadata\":"
+        << (gpu_metadata ? "true" : "false") << ','
+        << "\"timing_scope\":\"world_max\","
+        << "\"batch\":" << batch << ",\"seq\":" << seq << ','
+        << "\"hidden\":" << hidden << ",\"heads\":" << heads << ','
+        << "\"kv_heads\":" << kv_heads << ",\"head_dim\":" << head_dim << ','
+        << "\"layers\":" << layers << ",\"experts\":" << experts << ','
+        << "\"intermediate\":" << intermediate << ','
+        << "\"global_intermediate\":" << intermediate << ','
+        << "\"local_intermediate\":"
+        << (expert_tp ? intermediate / kTpSize : intermediate) << ','
+        << "\"expert_base_replication_factor\":"
+        << (expert_tp ? 1 : kTpSize) << ','
+        << "\"top_k\":" << top_k << ','
+        << "\"vocab\":" << vocab << ",\"lora_rank\":" << lora_rank << ','
+        << "\"heterogeneous_lora_rank\":" << heterogeneous_rank << ','
+        << "\"heterogeneous_targets\":\""
+        << json_escape(heterogeneous_targets) << "\","
+        << "\"tenants\":" << (dynamic_lora ? tenants : 0) << ','
+        << "\"active_tenants\":"
+        << (dynamic_lora ? active_tenants : 0) << ','
+        << "\"rotate_tenants\":" << (rotate_tenants ? "true" : "false") << ','
+        << "\"warmup\":" << warmup << ",\"iters\":" << iters << ','
+        << "\"last_loss\":" << last_loss << ','
+        << "\"step_ms_mean\":" << mean << ','
+        << "\"step_ms_p50\":" << p50 << ','
+        << "\"step_ms_p95\":" << p95 << ','
+        << "\"step_ms_std\":" << std::sqrt(variance) << ','
+        << "\"unique_tokens\":" << unique_tokens << ','
+        << "\"loss_tokens\":" << loss_tokens << ','
+        << "\"routed_tokens\":" << routed_tokens << ','
+        << "\"unique_tokens_per_sec\":" << unique_tokens_per_sec << ','
+        << "\"routed_tokens_per_sec\":" << routed_tokens_per_sec << ','
+        << "\"device_total_gib\":" << gib(total_bytes) << ','
+        << "\"free_start_gib\":" << gib(free_start) << ','
+        << "\"free_after_context_gib\":" << gib(free_after_context) << ','
+        << "\"free_after_warmup_gib\":" << gib(free_after_warmup) << ','
+        << "\"max_observed_resident_gib\":"
+        << gib(static_cast<int64_t>(used_since(free_start, min_observed_free))) << ','
+        << "\"allocator_current_allocated_gib\":" << gib(allocated.current) << ','
+        << "\"allocator_peak_allocated_gib\":" << gib(allocated.peak) << ','
+        << "\"allocator_current_reserved_gib\":" << gib(reserved.current) << ','
+        << "\"allocator_peak_reserved_gib\":" << gib(reserved.peak) << ','
+        << "\"samples_ms\":[";
+    for (size_t index = 0; index < times_ms.size(); ++index) {
+        if (index) output << ',';
+        output << times_ms[index];
+    }
+    output << "]}\n";
+    const std::string line = output.str();
+    std::fwrite(line.data(), 1, line.size(), stdout);
+    std::fflush(stdout);
+
+    qwen36_free_training_context(context);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_tp_ep_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_ep_smoke.cpp
new file mode 100644
index 00000000..ca0477b6
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_tp_ep_smoke.cpp
@@ -0,0 +1,1701 @@
+#include <ATen/ATen.h>
+
+#include <algorithm>
+#include <array>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <initializer_list>
+#include <limits>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" int64_t qwen36_kernel_abi_version();
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_parallel_nccl(
+    void*, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" int32_t qwen36_attach_parallel_nccl_no_sync(
+    void*, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" int32_t qwen36_set_mtp_weights(
+    void*, void*, void*, void*, void*, void**, int64_t, void*, int64_t);
+extern "C" int64_t qwen36_get_lora_count(void*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" int64_t qwen36_export_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" int64_t qwen36_import_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" int64_t qwen36_get_step_count(void*);
+extern "C" int32_t qwen36_set_step_count(void*, int64_t);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_micro_step(
+    void*, void*, void*, void*, double, int32_t);
+extern "C" int32_t qwen36_abort_gradient_accumulation(void*);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int64_t qwen36_add_lora_v2(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int64_t qwen36_add_lora_for_restore(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" void* qwen36_get_adapter_optimizer_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t);
+extern "C" int32_t qwen36_export_adapter_optimizer_tensor_cpu_v1(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t, void**);
+extern "C" int32_t qwen36_import_adapter_optimizer_state_host_v1(
+    void*, int64_t, const int64_t*, const char* const*, void* const*,
+    int64_t, int64_t);
+extern "C" int64_t qwen36_get_adapter_optimizer_resident_count_v1(
+    void*, int64_t);
+extern "C" int64_t qwen36_get_adapter_optimizer_resident_bytes_v1(
+    void*, int64_t);
+extern "C" int32_t qwen36_set_adapter_optimizer_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t, void*);
+extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t);
+extern "C" int32_t qwen36_get_context_health(void*);
+extern "C" int32_t qwen36_validate_adapter_steps_v1(
+    void*, const int64_t*, const int64_t*, int32_t);
+extern "C" int64_t qwen36_get_dynamic_finalizer_count(void*);
+extern "C" int64_t qwen36_get_dynamic_adam_launch_count(void*);
+extern "C" int64_t qwen36_get_dynamic_train_batch_count(void*);
+extern "C" int32_t qwen36_set_adapter_step_count(
+    void*, int64_t, int64_t);
+extern "C" double qwen36_train_multi_lora_selected(
+    void*, void*, void*, void*, const int64_t*, int32_t, int32_t);
+extern "C" double qwen36_train_multi_lora_selected_v2(
+    void*, void*, void*, void*, const int64_t*, int32_t);
+extern "C" int32_t qwen36_train_multi_lora_selected_v3(
+    void*, void*, void*, void*, const int64_t*, int32_t,
+    double*, double*, int32_t);
+extern "C" void qwen36_free_training_context(void*);
+extern "C" void qwen36_free_tensor(void*);
+
+namespace {
+
+constexpr int64_t kAbiVersion = 31;
+constexpr int32_t kBaseTpAttention = 1 << 0;
+constexpr int32_t kDataParallel = 1 << 1;
+constexpr int32_t kVocabParallel = 1 << 2;
+constexpr int32_t kExpertParallel = 1 << 3;
+constexpr int32_t kBaseTpMlp = 1 << 4;
+constexpr int64_t kHidden = 16;
+constexpr int64_t kVocab = 16;
+constexpr int64_t kHeads = 2;
+constexpr int64_t kKvHeads = 2;
+constexpr int64_t kHeadDim = 8;
+constexpr int64_t kExperts = 2;
+constexpr int64_t kIntermediate = 8;
+constexpr int64_t kLoraRank = 4;
+constexpr int64_t kLoraPairs = 9;
+constexpr int64_t kOptimizerSlots = 2 * kLoraPairs;
+constexpr double kLearningRate = 1e-3;
+constexpr double kBeta1 = 0.9;
+constexpr double kBeta2 = 0.999;
+constexpr double kAdamEps = 1e-8;
+
+int required_env_int(const char* name) {
+    const char* value = std::getenv(name);
+    assert(value && value[0] != '\0');
+    return std::atoi(value);
+}
+
+at::Tensor fingerprint(
+    std::initializer_list<int64_t> shape, double scale, int64_t offset
+) {
+    int64_t count = 1;
+    for (int64_t dim : shape) count *= dim;
+    auto values = at::arange(
+        count, at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    return ((values.add(offset).remainder(37) - 18.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+at::Tensor unit(std::initializer_list<int64_t> shape) {
+    return at::ones(
+        shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+}
+
+std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) {
+    assert(lhs.sizes() == rhs.sizes());
+    return (lhs.to(at::kFloat) - rhs.to(at::kFloat))
+        .abs().max().item<double>();
+}
+
+double update_norm(const at::Tensor& after, const at::Tensor& before) {
+    return (after.to(at::kFloat) - before.to(at::kFloat))
+        .abs().sum().item<double>();
+}
+
+at::Tensor adam_expected(
+    const at::Tensor& before, const at::Tensor& m, const at::Tensor& v
+) {
+    auto m_hat = m.to(at::kFloat) / (1.0 - kBeta1);
+    auto v_hat = v.to(at::kFloat) / (1.0 - kBeta2);
+    return (before.to(at::kFloat) -
+        kLearningRate * m_hat / (v_hat.sqrt() + kAdamEps))
+        .to(before.scalar_type());
+}
+
+struct Batch {
+    at::Tensor input_ids;
+    at::Tensor target_mask;
+    at::Tensor attention_mask;
+};
+
+Batch source_batch(int source_rank) {
+    auto long_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kLong);
+    auto float_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat);
+    auto bool_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kBool);
+    if (source_rank == 0) {
+        return {
+            at::tensor({1, 2, 3, 4}, long_opts).reshape({1, 4}),
+            at::tensor({0.0, 1.0, 0.0, 0.0}, float_opts).reshape({1, 4}),
+            at::ones({1, 4}, bool_opts),
+        };
+    }
+    if (source_rank == 1) {
+        return {
+            at::tensor({5, 6, 7, 8}, long_opts).reshape({1, 4}),
+            at::tensor({0.0, 1.0, 1.0, 1.0}, float_opts).reshape({1, 4}),
+            at::ones({1, 4}, bool_opts),
+        };
+    }
+    if (source_rank == 2) {
+        return {
+            at::tensor({9, 10, 11, 12}, long_opts).reshape({1, 4}),
+            at::tensor({0.0, 1.0, 0.0, 0.0}, float_opts).reshape({1, 4}),
+            at::ones({1, 4}, bool_opts),
+        };
+    }
+    assert(source_rank == 3);
+    return {
+        at::tensor({13, 14, 15, 1}, long_opts).reshape({1, 4}),
+        at::tensor({0.0, 1.0, 1.0, 0.0}, float_opts).reshape({1, 4}),
+        at::ones({1, 4}, bool_opts),
+    };
+}
+
+Batch full_batch(int source_count) {
+    auto first = source_batch(0);
+    auto second = source_batch(1);
+    if (source_count == 2) {
+        return {
+            at::cat({first.input_ids, second.input_ids}, 0),
+            at::cat({first.target_mask, second.target_mask}, 0),
+            at::cat({first.attention_mask, second.attention_mask}, 0),
+        };
+    }
+    assert(source_count == 4);
+    auto third = source_batch(2);
+    auto fourth = source_batch(3);
+    return {
+        at::cat({first.input_ids, second.input_ids,
+            third.input_ids, fourth.input_ids}, 0),
+        at::cat({first.target_mask, second.target_mask,
+            third.target_mask, fourth.target_mask}, 0),
+        at::cat({first.attention_mask, second.attention_mask,
+            third.attention_mask, fourth.attention_mask}, 0),
+    };
+}
+
+std::vector<at::Tensor> make_full_weights() {
+    std::vector<at::Tensor> weights;
+    weights.reserve(15);
+    weights.push_back(unit({kHidden}));
+    weights.push_back(unit({kHidden}));
+    weights.push_back(fingerprint(
+        {2 * kHeads * kHeadDim, kHidden}, 0.0011, 11));
+    weights.push_back(unit({kHeadDim}));
+    weights.push_back(fingerprint(
+        {kKvHeads * kHeadDim, kHidden}, 0.0012, 101));
+    weights.push_back(unit({kHeadDim}));
+    weights.push_back(fingerprint(
+        {kKvHeads * kHeadDim, kHidden}, 0.0013, 211));
+    weights.push_back(fingerprint(
+        {kHidden, kHeads * kHeadDim}, 0.0010, 307));
+    weights.push_back(fingerprint({kExperts, kHidden}, 0.0020, 401));
+    weights.push_back(fingerprint({1, kHidden}, 0.0015, 503));
+    weights.push_back(fingerprint({kIntermediate, kHidden}, 0.0010, 601));
+    weights.push_back(fingerprint({kIntermediate, kHidden}, 0.0011, 701));
+    weights.push_back(fingerprint({kHidden, kIntermediate}, 0.0012, 809));
+    weights.push_back(fingerprint(
+        {kExperts, 2 * kIntermediate, kHidden}, 0.0010, 907));
+    weights.push_back(fingerprint(
+        {kExperts, kHidden, kIntermediate}, 0.0011, 1009));
+    for (auto& weight : weights) weight.set_requires_grad(false);
+    return weights;
+}
+
+std::vector<at::Tensor> make_local_weights(
+    const std::vector<at::Tensor>& full, int tp_rank, int ep_rank,
+    bool base_tp_mlp
+) {
+    const int64_t local_heads = kHeads / 2;
+    const int64_t local_kv_heads = kKvHeads / 2;
+    const int64_t local_intermediate = kIntermediate / 2;
+    std::vector<at::Tensor> local;
+    local.reserve(full.size());
+    local.push_back(full[0]);
+    local.push_back(full[1]);
+    local.push_back(full[2].narrow(
+        0, tp_rank * 2 * local_heads * kHeadDim,
+        2 * local_heads * kHeadDim).contiguous());
+    local.push_back(full[3]);
+    local.push_back(full[4].narrow(
+        0, tp_rank * local_kv_heads * kHeadDim,
+        local_kv_heads * kHeadDim).contiguous());
+    local.push_back(full[5]);
+    local.push_back(full[6].narrow(
+        0, tp_rank * local_kv_heads * kHeadDim,
+        local_kv_heads * kHeadDim).contiguous());
+    local.push_back(full[7].narrow(
+        1, tp_rank * local_heads * kHeadDim,
+        local_heads * kHeadDim).contiguous());
+    local.push_back(full[8]);
+    local.push_back(full[9]);
+    local.push_back(base_tp_mlp ? full[10].narrow(
+        0, tp_rank * local_intermediate, local_intermediate).contiguous()
+        : full[10]);
+    local.push_back(base_tp_mlp ? full[11].narrow(
+        0, tp_rank * local_intermediate, local_intermediate).contiguous()
+        : full[11]);
+    local.push_back(base_tp_mlp ? full[12].narrow(
+        1, tp_rank * local_intermediate, local_intermediate).contiguous()
+        : full[12]);
+    auto local_gate_up = full[13].narrow(0, ep_rank, 1);
+    local.push_back(base_tp_mlp ? at::cat({
+            local_gate_up.narrow(
+                1, tp_rank * local_intermediate, local_intermediate),
+            local_gate_up.narrow(
+                1, kIntermediate + tp_rank * local_intermediate,
+                local_intermediate),
+        }, 1).contiguous()
+        : local_gate_up.contiguous());
+    auto local_down = full[14].narrow(0, ep_rank, 1);
+    local.push_back(base_tp_mlp ? local_down
+        .narrow(2, tp_rank * local_intermediate, local_intermediate)
+        .contiguous() : local_down.contiguous());
+    const int64_t expected_intermediate =
+        base_tp_mlp ? local_intermediate : kIntermediate;
+    assert(local[10].sizes() == at::IntArrayRef({expected_intermediate, kHidden}));
+    assert(local[11].sizes() == at::IntArrayRef({expected_intermediate, kHidden}));
+    assert(local[12].sizes() == at::IntArrayRef({kHidden, expected_intermediate}));
+    assert(local[13].sizes() == at::IntArrayRef(
+        {1, 2 * expected_intermediate, kHidden}));
+    assert(local[14].sizes() == at::IntArrayRef(
+        {1, kHidden, expected_intermediate}));
+    return local;
+}
+
+LayerConfig make_config(int expert_start, int expert_count) {
+    LayerConfig config{};
+    config.layer_type = 0;
+    config.num_heads = kHeads;
+    config.num_kv_heads = kKvHeads;
+    config.head_dim = kHeadDim;
+    config.partial_rotary_factor = 1.0;
+    config.rope_theta = 10000.0;
+    config.rms_eps = 1e-5;
+    config.num_experts = kExperts;
+    // top_k=2 makes every source exercise both owner ranks and every expert
+    // optimizer slot, independent of small-model router tie behavior.
+    config.top_k = 2;
+    config.moe_intermediate = kIntermediate;
+    config.expert_start = expert_start;
+    config.expert_count = expert_count;
+    config.norm_topk_prob = 1;
+    return config;
+}
+
+enum class FixtureKind {
+    Q,
+    SharedGate,
+    SharedDown,
+    ExpertGateUp,
+    ExpertDown,
+};
+
+struct LoraFixture {
+    int64_t slot;
+    const char* module;
+    FixtureKind kind;
+    at::Tensor full_a;
+    at::Tensor full_b;
+    at::Tensor local_a;
+    at::Tensor local_b;
+};
+
+at::Tensor local_factor(
+    const at::Tensor& full, FixtureKind kind, bool is_b,
+    int tp_rank, int ep_rank, bool base_tp_mlp
+) {
+    if (kind == FixtureKind::Q) {
+        if (!is_b) return full.clone();
+        const int64_t local_rows = full.size(0) / 2;
+        return full.narrow(0, tp_rank * local_rows, local_rows).contiguous();
+    }
+    const int64_t local_intermediate = kIntermediate / 2;
+    if (!base_tp_mlp) {
+        if (kind == FixtureKind::SharedGate ||
+            kind == FixtureKind::SharedDown) {
+            const int64_t rank_dim = is_b ? 1 : 0;
+            const int64_t local_rank = full.size(rank_dim) / 2;
+            return full.narrow(
+                rank_dim, tp_rank * local_rank, local_rank).contiguous();
+        }
+        auto local = full.narrow(0, ep_rank, 1);
+        const int64_t rank_dim = is_b ? 2 : 1;
+        const int64_t local_rank = local.size(rank_dim) / 2;
+        return local.narrow(
+            rank_dim, tp_rank * local_rank, local_rank).contiguous();
+    }
+    if (kind == FixtureKind::SharedGate) {
+        if (!is_b) return full.clone();
+        return full.narrow(
+            0, tp_rank * local_intermediate, local_intermediate).contiguous();
+    }
+    if (kind == FixtureKind::SharedDown) {
+        if (is_b) return full.clone();
+        return full.narrow(
+            1, tp_rank * local_intermediate, local_intermediate).contiguous();
+    }
+    auto local = full.narrow(0, ep_rank, 1);
+    if (kind == FixtureKind::ExpertGateUp) {
+        if (!is_b) return local.contiguous();
+        return at::cat({
+            local.narrow(
+                1, tp_rank * local_intermediate, local_intermediate),
+            local.narrow(
+                1, kIntermediate + tp_rank * local_intermediate,
+                local_intermediate),
+        }, 1).contiguous();
+    }
+    if (!is_b) {
+        return local.narrow(
+            2, tp_rank * local_intermediate, local_intermediate).contiguous();
+    }
+    return local.contiguous();
+}
+
+std::vector<LoraFixture> make_lora_fixtures(
+    int tp_rank, int ep_rank, int64_t offset, bool base_tp_mlp
+) {
+    std::vector<LoraFixture> result;
+    auto add = [&](int64_t slot, const char* module, FixtureKind kind,
+                   at::Tensor full_a, at::Tensor full_b) {
+        // Some narrow slices are already contiguous, so contiguous() may keep
+        // aliasing the full-reference tensor. DP perturbations must only touch
+        // the local fixture.
+        auto local_a = local_factor(
+            full_a, kind, false, tp_rank, ep_rank, base_tp_mlp).clone();
+        auto local_b = local_factor(
+            full_b, kind, true, tp_rank, ep_rank, base_tp_mlp).clone();
+        result.push_back({slot, module, kind, std::move(full_a),
+            std::move(full_b), std::move(local_a), std::move(local_b)});
+    };
+    add(0, "q_proj", FixtureKind::Q,
+        fingerprint({kLoraRank, kHidden}, 0.0007, offset + 1),
+        fingerprint({2 * kHeads * kHeadDim, kLoraRank},
+            0.0006, offset + 101));
+    add(4, "shared_gate_proj", FixtureKind::SharedGate,
+        fingerprint({kLoraRank, kHidden}, 0.0007, offset + 151),
+        fingerprint({kIntermediate, kLoraRank}, 0.0006, offset + 181));
+    add(6, "shared_down_proj", FixtureKind::SharedDown,
+        fingerprint({kLoraRank, kIntermediate}, 0.0007, offset + 191),
+        fingerprint({kHidden, kLoraRank}, 0.0006, offset + 201));
+    add(7, "experts_gate_up_proj", FixtureKind::ExpertGateUp,
+        fingerprint({kExperts, kLoraRank, kHidden}, 0.0007, offset + 211),
+        fingerprint({kExperts, 2 * kIntermediate, kLoraRank},
+            0.0006, offset + 307));
+    add(8, "experts_down_proj", FixtureKind::ExpertDown,
+        fingerprint({kExperts, kLoraRank, kIntermediate},
+            0.0007, offset + 401),
+        fingerprint({kExperts, kHidden, kLoraRank},
+            0.0006, offset + 503));
+    assert(result[0].local_a.sizes() == at::IntArrayRef({kLoraRank, kHidden}));
+    assert(result[0].local_b.sizes() == at::IntArrayRef({kHidden, kLoraRank}));
+    const int64_t local_rank = base_tp_mlp ? kLoraRank : kLoraRank / 2;
+    const int64_t local_intermediate =
+        base_tp_mlp ? kIntermediate / 2 : kIntermediate;
+    assert(result[1].local_a.sizes() == at::IntArrayRef({local_rank, kHidden}));
+    assert(result[1].local_b.sizes() == at::IntArrayRef({local_intermediate, local_rank}));
+    assert(result[2].local_a.sizes() == at::IntArrayRef({local_rank, local_intermediate}));
+    assert(result[2].local_b.sizes() == at::IntArrayRef({kHidden, local_rank}));
+    assert(result[3].local_a.sizes() == at::IntArrayRef({1, local_rank, kHidden}));
+    const int64_t local_gate_up_rows =
+        base_tp_mlp ? kIntermediate : 2 * kIntermediate;
+    assert(result[3].local_b.sizes() ==
+        at::IntArrayRef({1, local_gate_up_rows, local_rank}));
+    assert(result[4].local_a.sizes() == at::IntArrayRef({1, local_rank, local_intermediate}));
+    assert(result[4].local_b.sizes() == at::IntArrayRef({1, kHidden, local_rank}));
+    return result;
+}
+
+void install_fixed(
+    void* context, std::vector<LoraFixture>& fixtures, bool local
+) {
+    assert(qwen36_get_lora_count(context) == kLoraPairs);
+    for (auto& fixture : fixtures) {
+        auto& a = local ? fixture.local_a : fixture.full_a;
+        auto& b = local ? fixture.local_b : fixture.full_b;
+        assert(qwen36_set_lora_tensor(context, fixture.slot, 0, &a) == 0);
+        assert(qwen36_set_lora_tensor(context, fixture.slot, 1, &b) == 0);
+    }
+}
+
+void install_dynamic(
+    void* context, int64_t adapter_id,
+    std::vector<LoraFixture>& fixtures, bool local
+) {
+    for (auto& fixture : fixtures) {
+        auto& a = local ? fixture.local_a : fixture.full_a;
+        auto& b = local ? fixture.local_b : fixture.full_b;
+        assert(qwen36_set_adapter_lora_tensor(
+            context, adapter_id, 0, fixture.module, 0, &a) == 0);
+        assert(qwen36_set_adapter_lora_tensor(
+            context, adapter_id, 0, fixture.module, 1, &b) == 0);
+    }
+}
+
+at::Tensor* dynamic_tensor(
+    void* context, int64_t adapter_id, const char* module, bool is_b
+) {
+    auto* tensor = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_lora_tensor(
+            context, adapter_id, 0, module, is_b ? 1 : 0));
+    assert(tensor);
+    return tensor;
+}
+
+at::Tensor* dynamic_state(
+    void* context, int64_t adapter_id, const char* module,
+    bool is_b, bool is_v
+) {
+    auto* tensor = reinterpret_cast<at::Tensor*>(
+        qwen36_get_adapter_optimizer_tensor(
+            context, adapter_id, 0, module,
+            is_b ? 1 : 0, is_v ? 1 : 0));
+    assert(tensor && tensor->scalar_type() == at::kFloat);
+    return tensor;
+}
+
+at::Tensor dynamic_state_cpu(
+    void* context, int64_t adapter_id, const char* module,
+    bool is_b, bool is_v
+) {
+    void* output = nullptr;
+    assert(qwen36_export_adapter_optimizer_tensor_cpu_v1(
+        context, adapter_id, 0, module,
+        is_b ? 1 : 0, is_v ? 1 : 0, &output) == 0);
+    auto* tensor = reinterpret_cast<at::Tensor*>(output);
+    assert(tensor && tensor->device().is_cpu() &&
+        tensor->scalar_type() == at::kFloat && tensor->is_contiguous());
+    auto result = *tensor;
+    qwen36_free_tensor(output);
+    return result;
+}
+
+void set_distributed_env(
+    int rank, int world, int tp_rank, int ep_rank, int dp_rank, int dp_size
+) {
+    setenv("WORLD_SIZE", std::to_string(world).c_str(), 1);
+    setenv("TP_SIZE", "2", 1);
+    setenv("EP_SIZE", "2", 1);
+    setenv("DP_SIZE", std::to_string(dp_size).c_str(), 1);
+    setenv("RUSTRAIN_TP_RANK", std::to_string(tp_rank).c_str(), 1);
+    setenv("RUSTRAIN_EP_RANK", std::to_string(ep_rank).c_str(), 1);
+    setenv("RUSTRAIN_DP_RANK", std::to_string(dp_rank).c_str(), 1);
+    setenv("RUSTRAIN_DATA_PARALLEL", dp_size > 1 ? "1" : "0", 1);
+    setenv("RANK", std::to_string(rank).c_str(), 1);
+}
+
+void set_reference_env() {
+    setenv("WORLD_SIZE", "1", 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    setenv("RUSTRAIN_DATA_PARALLEL", "0", 1);
+    setenv("RANK", "0", 1);
+}
+
+struct ParityErrors {
+    double param = 0.0;
+    double m = 0.0;
+    double v = 0.0;
+    double adam = 0.0;
+};
+
+}  // namespace
+
+int main() {
+    const int rank = required_env_int("RANK");
+    const int world = required_env_int("WORLD_SIZE");
+    const int local_rank = required_env_int("LOCAL_RANK");
+    assert((world == 4 || world == 8) && rank >= 0 && rank < world);
+    const int dp_size = world / 4;
+    const int tp_rank = rank % 2;
+    const int ep_rank = (rank / 2) % 2;
+    const int dp_rank = rank / 4;
+    const int tp_color = ep_rank + 2 * dp_rank;
+    const int ep_color = tp_rank + 2 * dp_rank;
+    const int dp_color = tp_rank + 2 * ep_rank;
+    const char* sharded_a2a_env = std::getenv("QWEN36_EP_A2A_SHARDED");
+    assert(sharded_a2a_env);
+    const bool sharded_source = std::strcmp(sharded_a2a_env, "0") != 0;
+    const bool base_tp_mlp = sharded_source;
+    const int32_t distributed_flags =
+        kBaseTpAttention | kVocabParallel | kExpertParallel |
+        (base_tp_mlp ? kBaseTpMlp : 0) |
+        (dp_size > 1 ? kDataParallel : 0);
+    assert(qwen36_kernel_abi_version() == kAbiVersion);
+    assert(std::getenv("QWEN36_EP_A2A") &&
+        std::strcmp(std::getenv("QWEN36_EP_A2A"), "0") != 0);
+    qwen36_set_cuda_device(local_rank);
+
+    auto full_weights = make_full_weights();
+    auto local_weights = make_local_weights(
+        full_weights, tp_rank, ep_rank, base_tp_mlp);
+    auto full_ptrs = pointers(full_weights);
+    auto local_ptrs = pointers(local_weights);
+    auto embed = fingerprint({kVocab, kHidden}, 0.0012, 1201);
+    auto final_norm = unit({kHidden});
+    auto lm_head = fingerprint({kVocab, kHidden}, 0.0010, 1301);
+    auto local_embed = embed.narrow(
+        0, tp_rank * (kVocab / 2), kVocab / 2).contiguous();
+    auto local_lm_head = lm_head.narrow(
+        0, tp_rank * (kVocab / 2), kVocab / 2).contiguous();
+    embed.set_requires_grad(false);
+    final_norm.set_requires_grad(false);
+    lm_head.set_requires_grad(false);
+
+    auto distributed_config = make_config(ep_rank, 1);
+    auto reference_config = make_config(0, kExperts);
+    const int64_t target_layer = 0;
+    constexpr const char* targets =
+        "q_proj,shared_gate_proj,shared_down_proj,"
+        "experts_gate_up_proj,experts_down_proj";
+    constexpr const char* projection_targets = "q_proj";
+
+    // Reject a process grid that cannot cover WORLD_SIZE before any NCCL
+    // communicator is created.
+    set_distributed_env(rank, world, tp_rank, ep_rank, dp_rank, dp_size);
+    setenv("EP_SIZE", "3", 1);
+    void* invalid = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head, &distributed_config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        kVocab, 1e-5, kLoraRank, &target_layer, 1, targets,
+        distributed_flags);
+    assert(invalid == nullptr);
+    set_distributed_env(rank, world, tp_rank, ep_rank, dp_rank, dp_size);
+
+    void* projection_rank_three = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head, &distributed_config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        kVocab, 1e-5, 3, &target_layer, 1, projection_targets,
+        distributed_flags);
+    assert(projection_rank_three);
+    assert(qwen36_add_lora(
+        projection_rank_three, 3, 3.0, &target_layer, 1,
+        projection_targets) > 0);
+    qwen36_free_training_context(projection_rank_three);
+
+    constexpr const char* mixed_targets = "q_proj,shared_gate_proj";
+    const int32_t mixed_flags = distributed_flags & ~kBaseTpMlp;
+    void* invalid_mixed_rank_three = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head, &distributed_config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        kVocab, 1e-5, 3, &target_layer, 1, mixed_targets,
+        mixed_flags);
+    assert(invalid_mixed_rank_three == nullptr);
+    void* mixed_rank_four = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head, &distributed_config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        kVocab, 1e-5, kLoraRank, &target_layer, 1, mixed_targets,
+        mixed_flags);
+    assert(mixed_rank_four);
+    assert(qwen36_add_lora(
+        mixed_rank_four, 3, 3.0, &target_layer, 1, mixed_targets) == -1);
+    qwen36_free_training_context(mixed_rank_four);
+
+    void* distributed = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head, &distributed_config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        kVocab, 1e-5, kLoraRank, &target_layer, 1, targets,
+        distributed_flags);
+    assert(distributed);
+
+    set_reference_env();
+    void* reference = qwen36_create_training_context(
+        full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head,
+        &reference_config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        kVocab, 1e-5, kLoraRank, &target_layer, 1, targets);
+    assert(reference);
+    set_distributed_env(rank, world, tp_rank, ep_rank, dp_rank, dp_size);
+
+    // Frozen-base/vocabulary TP must reject unsharded MTP weights directly.
+    auto mtp_dummy = unit({1});
+    assert(qwen36_set_mtp_weights(
+        distributed, &mtp_dummy, &mtp_dummy, &mtp_dummy, &mtp_dummy,
+        nullptr, 0, nullptr, 0) == -1);
+
+    auto fixed_fixtures = make_lora_fixtures(
+        tp_rank, ep_rank, 2001, base_tp_mlp);
+    if (dp_rank > 0) {
+        for (auto& fixture : fixed_fixtures) {
+            fixture.local_a.add_(0.25);
+            fixture.local_b.add_(0.25);
+        }
+    }
+    install_fixed(distributed, fixed_fixtures, true);
+    install_fixed(reference, fixed_fixtures, false);
+
+    // Default tp-ep-dp rank order makes TP the least-significant coordinate.
+    assert(qwen36_init_parallel_nccl(
+        distributed, rank, world,
+        0, 1, 0,
+        rank % 2, 2, rank / 2,
+        rank / 2, 2, rank % 2) == -1);
+    assert(qwen36_init_parallel_nccl(
+        distributed, rank, world,
+        tp_rank, 2, tp_color,
+        ep_rank, 2, ep_color,
+        dp_rank, dp_size, dp_color) == 0);
+
+    std::vector<std::array<at::Tensor, 2>> fixed_before;
+    double broadcast_diff = 0.0;
+    for (auto& fixture : fixed_fixtures) {
+        auto* a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(distributed, fixture.slot));
+        auto* b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(distributed, fixture.slot));
+        assert(a && b);
+        auto expected_a = local_factor(
+            fixture.full_a, fixture.kind, false, tp_rank, ep_rank,
+            base_tp_mlp);
+        auto expected_b = local_factor(
+            fixture.full_b, fixture.kind, true, tp_rank, ep_rank,
+            base_tp_mlp);
+        const double a_diff = max_diff(*a, expected_a);
+        const double b_diff = max_diff(*b, expected_b);
+        if (a_diff != 0.0 || b_diff != 0.0) {
+            std::fprintf(stderr,
+                "fixed_broadcast_mismatch rank=%d tp=%d ep=%d dp=%d "
+                "module=%s a_diff=%0.8e b_diff=%0.8e\n",
+                rank, tp_rank, ep_rank, dp_rank, fixture.module,
+                a_diff, b_diff);
+        }
+        broadcast_diff = std::max({broadcast_diff,
+            a_diff, b_diff});
+        fixed_before.push_back({a->clone(), b->clone()});
+    }
+    assert(broadcast_diff == 0.0);
+
+    const int source_rank = sharded_source ? 2 * dp_rank + ep_rank : dp_rank;
+    auto local_batch = source_batch(source_rank);
+    const int source_count = sharded_source ? 2 * dp_size : dp_size;
+    auto global_batch = full_batch(source_count);
+
+    auto* fixed_preflight_probe = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_b(distributed, fixed_fixtures.front().slot));
+    assert(fixed_preflight_probe);
+    const auto fixed_preflight_probe_before = fixed_preflight_probe->clone();
+
+    // Replica input signatures must agree before TP/EP forward collectives.
+    auto mismatched_shape_input = local_batch.input_ids;
+    auto mismatched_shape_targets = local_batch.target_mask;
+    auto mismatched_shape_attention = local_batch.attention_mask;
+    if (rank == 0) {
+        mismatched_shape_input = mismatched_shape_input.narrow(1, 0, 3).contiguous();
+        mismatched_shape_targets = mismatched_shape_targets.narrow(1, 0, 3).contiguous();
+        mismatched_shape_attention = mismatched_shape_attention.narrow(1, 0, 3).contiguous();
+    }
+    assert(qwen36_train_micro_step(
+        distributed, &mismatched_shape_input, &mismatched_shape_targets,
+        &mismatched_shape_attention, 1.0, 1) < 0.0);
+    assert(qwen36_get_step_count(distributed) == 0);
+    assert(max_diff(*fixed_preflight_probe, fixed_preflight_probe_before) == 0.0);
+
+    // The scale itself is finite, but sources with multiple supervised tokens
+    // overflow the product. Every rank must fail before forward/backward.
+    assert(qwen36_train_micro_step(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask,
+        std::numeric_limits<double>::max(), 1) < 0.0);
+    assert(qwen36_get_step_count(distributed) == 0);
+    assert(max_diff(*fixed_preflight_probe, fixed_preflight_probe_before) == 0.0);
+
+    // Optimizer phase is part of the distributed fixed-LoRA clock. A rank-local
+    // apply decision must fail before forward or accumulation can diverge.
+    const double mismatched_phase_loss = qwen36_train_micro_step(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, 1.0, rank == 0 ? 0 : 1);
+    assert(mismatched_phase_loss < 0.0);
+    assert(qwen36_get_step_count(distributed) == 0);
+    assert(max_diff(*fixed_preflight_probe, fixed_preflight_probe_before) == 0.0);
+
+    // A rank-local fixed optimizer clock must fail through full-topology
+    // consensus before gradient collectives or Adam can diverge.
+    auto* fixed_clock_probe = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_b(distributed, fixed_fixtures.front().slot));
+    assert(fixed_clock_probe);
+    const auto fixed_clock_probe_before = fixed_clock_probe->clone();
+    if (rank == 0)
+        assert(qwen36_set_step_count(distributed, 1) == 0);
+    const double mismatched_clock_loss = qwen36_train_step(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask);
+    assert(mismatched_clock_loss < 0.0);
+    if (rank == 0)
+        assert(qwen36_set_step_count(distributed, 0) == 0);
+    assert(qwen36_get_step_count(distributed) == 0);
+    assert(max_diff(*fixed_clock_probe, fixed_clock_probe_before) == 0.0);
+
+    // Registry mutation is collective and must reject disagreement about a
+    // pending accumulation window before adapter parameter synchronization.
+    assert(qwen36_train_micro_step(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, 1.0, 0) > 0.0);
+    if (rank == 0)
+        assert(qwen36_abort_gradient_accumulation(distributed) == 0);
+    assert(qwen36_add_lora(
+        distributed, kLoraRank, 1.0, &target_layer, 1, targets) < 0);
+    assert(qwen36_abort_gradient_accumulation(distributed) == 0);
+    assert(qwen36_get_step_count(distributed) == 0);
+
+    const double distributed_loss = qwen36_train_step(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask);
+    const double reference_loss = qwen36_train_step(
+        reference, &global_batch.input_ids, &global_batch.target_mask,
+        &global_batch.attention_mask);
+    assert(distributed_loss > 0.0 && std::isfinite(distributed_loss));
+    assert(reference_loss > 0.0 && std::isfinite(reference_loss));
+
+    std::vector<void*> local_m(kOptimizerSlots), local_v(kOptimizerSlots);
+    std::vector<void*> full_m(kOptimizerSlots), full_v(kOptimizerSlots);
+    assert(qwen36_export_optimizer_state(
+        distributed, local_m.data(), local_v.data(), kOptimizerSlots) ==
+        kOptimizerSlots);
+    assert(qwen36_export_optimizer_state(
+        reference, full_m.data(), full_v.data(), kOptimizerSlots) ==
+        kOptimizerSlots);
+
+    ParityErrors fixed_errors;
+    for (size_t fixture_index = 0;
+         fixture_index < fixed_fixtures.size(); ++fixture_index) {
+        auto& fixture = fixed_fixtures[fixture_index];
+        auto* local_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(distributed, fixture.slot));
+        auto* local_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(distributed, fixture.slot));
+        auto* ref_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(reference, fixture.slot));
+        auto* ref_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(reference, fixture.slot));
+        assert(local_a && local_b && ref_a && ref_b);
+        fixed_errors.param = std::max({fixed_errors.param,
+            max_diff(*local_a, local_factor(
+                *ref_a, fixture.kind, false, tp_rank, ep_rank,
+                base_tp_mlp)),
+            max_diff(*local_b, local_factor(
+                *ref_b, fixture.kind, true, tp_rank, ep_rank,
+                base_tp_mlp))});
+
+        const int64_t a_state = 2 * fixture.slot;
+        const int64_t b_state = a_state + 1;
+        auto* m_a = reinterpret_cast<at::Tensor*>(local_m[a_state]);
+        auto* m_b = reinterpret_cast<at::Tensor*>(local_m[b_state]);
+        auto* v_a = reinterpret_cast<at::Tensor*>(local_v[a_state]);
+        auto* v_b = reinterpret_cast<at::Tensor*>(local_v[b_state]);
+        auto* ref_m_a = reinterpret_cast<at::Tensor*>(full_m[a_state]);
+        auto* ref_m_b = reinterpret_cast<at::Tensor*>(full_m[b_state]);
+        auto* ref_v_a = reinterpret_cast<at::Tensor*>(full_v[a_state]);
+        auto* ref_v_b = reinterpret_cast<at::Tensor*>(full_v[b_state]);
+        assert(m_a && m_b && v_a && v_b &&
+            ref_m_a && ref_m_b && ref_v_a && ref_v_b);
+        fixed_errors.m = std::max({fixed_errors.m,
+            max_diff(*m_a, local_factor(
+                *ref_m_a, fixture.kind, false, tp_rank, ep_rank,
+                base_tp_mlp)),
+            max_diff(*m_b, local_factor(
+                *ref_m_b, fixture.kind, true, tp_rank, ep_rank,
+                base_tp_mlp))});
+        fixed_errors.v = std::max({fixed_errors.v,
+            max_diff(*v_a, local_factor(
+                *ref_v_a, fixture.kind, false, tp_rank, ep_rank,
+                base_tp_mlp)),
+            max_diff(*v_b, local_factor(
+                *ref_v_b, fixture.kind, true, tp_rank, ep_rank,
+                base_tp_mlp))});
+        fixed_errors.adam = std::max({fixed_errors.adam,
+            max_diff(*local_a, adam_expected(
+                fixed_before[fixture_index][0], *m_a, *v_a)),
+            max_diff(*local_b, adam_expected(
+                fixed_before[fixture_index][1], *m_b, *v_b))});
+        assert(update_norm(*local_a, fixed_before[fixture_index][0]) > 0.0);
+        assert(update_norm(*local_b, fixed_before[fixture_index][1]) > 0.0);
+    }
+
+    std::printf(
+        "native_tp_ep_fixed rank=%d tp=%d ep=%d dp=%d source_mode=%s "
+        "local_loss=%0.8f "
+        "reference_loss=%0.8f param_diff=%0.8e m_diff=%0.8e "
+        "v_diff=%0.8e adam_error=%0.8e\n",
+        rank, tp_rank, ep_rank, dp_rank,
+        sharded_source ? "sharded" : "replicated",
+        distributed_loss, reference_loss,
+        fixed_errors.param, fixed_errors.m, fixed_errors.v,
+        fixed_errors.adam);
+    std::fflush(stdout);
+    // TP matmul reductions and EP dispatch round BF16 in a different order
+    // from the full model. FP32 optimizer states remain the tight oracle.
+    assert(fixed_errors.param <= 3e-3);
+    assert(fixed_errors.m <= 7e-3);
+    assert(fixed_errors.v <= 2e-5);
+    assert(fixed_errors.adam <= 1e-5);
+
+    // Simulate a safetensors checkpoint: fixed LoRA and Adam state leave the
+    // device, then a fresh distributed context restores them from CPU tensors.
+    std::vector<at::Tensor> checkpoint_a;
+    std::vector<at::Tensor> checkpoint_b;
+    std::vector<at::Tensor> checkpoint_m;
+    std::vector<at::Tensor> checkpoint_v;
+    checkpoint_a.reserve(kLoraPairs);
+    checkpoint_b.reserve(kLoraPairs);
+    checkpoint_m.reserve(kOptimizerSlots);
+    checkpoint_v.reserve(kOptimizerSlots);
+    for (int64_t slot = 0; slot < kLoraPairs; ++slot) {
+        auto* a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(distributed, slot));
+        auto* b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(distributed, slot));
+        assert(a && b);
+        checkpoint_a.push_back(a->to(at::kCPU).clone());
+        checkpoint_b.push_back(b->to(at::kCPU).clone());
+    }
+    for (int64_t index = 0; index < kOptimizerSlots; ++index) {
+        auto* m = reinterpret_cast<at::Tensor*>(local_m[index]);
+        auto* v = reinterpret_cast<at::Tensor*>(local_v[index]);
+        assert(m && v);
+        checkpoint_m.push_back(m->to(at::kCPU).clone());
+        checkpoint_v.push_back(v->to(at::kCPU).clone());
+    }
+
+    void* resumed = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head, &distributed_config, 1,
+        static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        kVocab, 1e-5, kLoraRank, &target_layer, 1, targets,
+        distributed_flags);
+    assert(resumed);
+    std::vector<std::array<at::Tensor, 2>> shadow_fixed_before_attach;
+    shadow_fixed_before_attach.reserve(kLoraPairs);
+    for (int64_t slot = 0; slot < kLoraPairs; ++slot) {
+        auto* a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(resumed, slot));
+        auto* b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(resumed, slot));
+        assert(a && b);
+        if (dp_rank > 0) {
+            auto divergent_a = a->detach().clone().add_(0.125);
+            auto divergent_b = b->detach().clone().add_(0.25);
+            assert(qwen36_set_lora_tensor(
+                resumed, slot, 0, &divergent_a) == 0);
+            assert(qwen36_set_lora_tensor(
+                resumed, slot, 1, &divergent_b) == 0);
+        }
+        shadow_fixed_before_attach.push_back({a->clone(), b->clone()});
+    }
+    assert(qwen36_attach_parallel_nccl_no_sync(
+        resumed, rank, world,
+        tp_rank, 2, tp_color,
+        ep_rank, 2, ep_color,
+        dp_rank, dp_size, dp_color) == 0);
+    for (int64_t slot = 0; slot < kLoraPairs; ++slot) {
+        auto* a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(resumed, slot));
+        auto* b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(resumed, slot));
+        assert(a && b);
+        assert(max_diff(*a, shadow_fixed_before_attach[slot][0]) == 0.0);
+        assert(max_diff(*b, shadow_fixed_before_attach[slot][1]) == 0.0);
+    }
+    for (int64_t slot = 0; slot < kLoraPairs; ++slot) {
+        assert(qwen36_set_lora_tensor(
+            resumed, slot, 0, &checkpoint_a[slot]) == 0);
+        assert(qwen36_set_lora_tensor(
+            resumed, slot, 1, &checkpoint_b[slot]) == 0);
+    }
+    auto checkpoint_m_ptrs = pointers(checkpoint_m);
+    auto checkpoint_v_ptrs = pointers(checkpoint_v);
+    assert(qwen36_import_optimizer_state(
+        resumed, checkpoint_m_ptrs.data(), checkpoint_v_ptrs.data(),
+        kOptimizerSlots) == kOptimizerSlots);
+    assert(qwen36_set_step_count(
+        resumed, qwen36_get_step_count(distributed)) == 0);
+
+    std::vector<void*> resumed_m(kOptimizerSlots), resumed_v(kOptimizerSlots);
+    assert(qwen36_export_optimizer_state(
+        resumed, resumed_m.data(), resumed_v.data(), kOptimizerSlots) ==
+        kOptimizerSlots);
+    double restored_param_diff = 0.0;
+    double restored_m_diff = 0.0;
+    double restored_v_diff = 0.0;
+    for (int64_t slot = 0; slot < kLoraPairs; ++slot) {
+        auto* original_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(distributed, slot));
+        auto* original_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(distributed, slot));
+        auto* restored_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(resumed, slot));
+        auto* restored_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(resumed, slot));
+        assert(original_a && original_b && restored_a && restored_b);
+        assert(restored_a->is_cuda() && restored_b->is_cuda());
+        restored_param_diff = std::max({restored_param_diff,
+            max_diff(*original_a, *restored_a),
+            max_diff(*original_b, *restored_b)});
+    }
+    for (int64_t index = 0; index < kOptimizerSlots; ++index) {
+        auto* original_m = reinterpret_cast<at::Tensor*>(local_m[index]);
+        auto* original_v = reinterpret_cast<at::Tensor*>(local_v[index]);
+        auto* restored_m = reinterpret_cast<at::Tensor*>(resumed_m[index]);
+        auto* restored_v = reinterpret_cast<at::Tensor*>(resumed_v[index]);
+        assert(original_m && original_v && restored_m && restored_v);
+        assert(restored_m->is_cuda() && restored_v->is_cuda());
+        assert(restored_m->scalar_type() == at::kFloat);
+        assert(restored_v->scalar_type() == at::kFloat);
+        restored_m_diff = std::max(
+            restored_m_diff, max_diff(*original_m, *restored_m));
+        restored_v_diff = std::max(
+            restored_v_diff, max_diff(*original_v, *restored_v));
+    }
+    assert(restored_param_diff == 0.0);
+    assert(restored_m_diff == 0.0);
+    assert(restored_v_diff == 0.0);
+    assert(qwen36_get_step_count(resumed) ==
+        qwen36_get_step_count(distributed));
+
+    const double continued_loss = qwen36_train_step(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask);
+    const double resumed_loss = qwen36_train_step(
+        resumed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask);
+    assert(std::abs(continued_loss - resumed_loss) <= 1e-6);
+    assert(qwen36_export_optimizer_state(
+        distributed, local_m.data(), local_v.data(), kOptimizerSlots) ==
+        kOptimizerSlots);
+    assert(qwen36_export_optimizer_state(
+        resumed, resumed_m.data(), resumed_v.data(), kOptimizerSlots) ==
+        kOptimizerSlots);
+    double continued_param_diff = 0.0;
+    double continued_m_diff = 0.0;
+    double continued_v_diff = 0.0;
+    for (int64_t slot = 0; slot < kLoraPairs; ++slot) {
+        auto* continued_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(distributed, slot));
+        auto* continued_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(distributed, slot));
+        auto* restored_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(resumed, slot));
+        auto* restored_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(resumed, slot));
+        continued_param_diff = std::max({continued_param_diff,
+            max_diff(*continued_a, *restored_a),
+            max_diff(*continued_b, *restored_b)});
+    }
+    for (int64_t index = 0; index < kOptimizerSlots; ++index) {
+        continued_m_diff = std::max(continued_m_diff, max_diff(
+            *reinterpret_cast<at::Tensor*>(local_m[index]),
+            *reinterpret_cast<at::Tensor*>(resumed_m[index])));
+        continued_v_diff = std::max(continued_v_diff, max_diff(
+            *reinterpret_cast<at::Tensor*>(local_v[index]),
+            *reinterpret_cast<at::Tensor*>(resumed_v[index])));
+    }
+    std::printf(
+        "native_tp_ep_resume rank=%d loss=%0.8f param_diff=%0.8e "
+        "m_diff=%0.8e v_diff=%0.8e step=%ld\n",
+        rank, resumed_loss, continued_param_diff, continued_m_diff,
+        continued_v_diff,
+        static_cast<long>(qwen36_get_step_count(resumed)));
+    std::fflush(stdout);
+    assert(continued_param_diff == 0.0);
+    assert(continued_m_diff == 0.0);
+    assert(continued_v_diff == 0.0);
+    assert(qwen36_get_step_count(resumed) == 2);
+    const int64_t fixed_step_before_dynamic =
+        qwen36_get_step_count(distributed);
+
+    const int64_t dynamic_targets[] = {0};
+    const int64_t divergent_rank = rank == 0 ? kLoraRank - 1 : kLoraRank;
+    assert(qwen36_add_lora(
+        distributed, divergent_rank, kLoraRank,
+        dynamic_targets, 1, targets) < 0);
+    const char* divergent_targets = rank == 0 ? "q_proj" : targets;
+    assert(qwen36_add_lora(
+        distributed, kLoraRank, kLoraRank,
+        dynamic_targets, 1, divergent_targets) < 0);
+    const int64_t divergent_mode_result = rank == 0
+        ? qwen36_add_lora(
+            distributed, kLoraRank, kLoraRank,
+            dynamic_targets, 1, targets)
+        : qwen36_add_lora_v2(
+            distributed, kLoraRank, kLoraRank,
+            dynamic_targets, 1, targets);
+    assert(divergent_mode_result < 0);
+    if (rank == 0) {
+        setenv(
+            "QWEN36_TEST_FAIL_ADAPTER_REGISTRATION_AFTER_SYNC", "1", 1);
+    }
+    assert(qwen36_add_lora(
+        distributed, kLoraRank, kLoraRank,
+        dynamic_targets, 1, targets) < 0);
+    unsetenv("QWEN36_TEST_FAIL_ADAPTER_REGISTRATION_AFTER_SYNC");
+    const int64_t tenant_one = qwen36_add_lora(
+        distributed, kLoraRank, kLoraRank,
+        dynamic_targets, 1, targets);
+    const int64_t tenant_two = qwen36_add_lora(
+        distributed, kLoraRank, kLoraRank,
+        dynamic_targets, 1, targets);
+    assert(tenant_one == 1 && tenant_two == 2);
+    auto dynamic_one = make_lora_fixtures(
+        tp_rank, ep_rank, 4001, base_tp_mlp);
+    auto dynamic_two = make_lora_fixtures(
+        tp_rank, ep_rank, 6001, base_tp_mlp);
+    auto reference_dynamic_one = make_lora_fixtures(
+        tp_rank, ep_rank, 4001, base_tp_mlp);
+    install_dynamic(distributed, tenant_one, dynamic_one, true);
+    install_dynamic(distributed, tenant_two, dynamic_two, true);
+    void* untrained_state = reinterpret_cast<void*>(1);
+    assert(qwen36_export_adapter_optimizer_tensor_cpu_v1(
+        distributed, tenant_one, 0, dynamic_one.front().module,
+        0, 0, &untrained_state) == 1);
+    assert(untrained_state == nullptr);
+    // A selected dynamic tenant owns one activation row locally. Its full-rank
+    // oracle must aggregate every source row into one adapter update, which a
+    // fresh fixed-LoRA context does without conflating rows with tenant IDs.
+    set_reference_env();
+    void* dynamic_reference = qwen36_create_training_context(
+        full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head,
+        &reference_config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, kLearningRate, kBeta1, kBeta2, kAdamEps,
+        kVocab, 1e-5, kLoraRank, &target_layer, 1, targets);
+    assert(dynamic_reference);
+    install_fixed(dynamic_reference, reference_dynamic_one, false);
+    set_distributed_env(rank, world, tp_rank, ep_rank, dp_rank, dp_size);
+    assert(qwen36_train_step(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask) < 0.0);
+    assert(qwen36_get_step_count(distributed) == fixed_step_before_dynamic);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) == 0);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_two) == 0);
+    const int64_t guarded_tenants[] = {tenant_one, tenant_two};
+    const int64_t initial_steps[] = {0, 0};
+    assert(qwen36_validate_adapter_steps_v1(
+        distributed, guarded_tenants, initial_steps, 2) == 0);
+    const int64_t divergent_expected_steps[] = {
+        rank == 0 ? 1 : 0, 0};
+    assert(qwen36_validate_adapter_steps_v1(
+        distributed, guarded_tenants, divergent_expected_steps, 2) < 0);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) == 0);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_two) == 0);
+
+    std::vector<std::array<at::Tensor, 2>> tenant_one_before;
+    std::vector<std::array<at::Tensor, 2>> tenant_two_before;
+    for (size_t index = 0; index < dynamic_one.size(); ++index) {
+        tenant_one_before.push_back({
+            dynamic_tensor(distributed, tenant_one,
+                dynamic_one[index].module, false)->clone(),
+            dynamic_tensor(distributed, tenant_one,
+                dynamic_one[index].module, true)->clone(),
+        });
+        tenant_two_before.push_back({
+            dynamic_tensor(distributed, tenant_two,
+                dynamic_two[index].module, false)->clone(),
+            dynamic_tensor(distributed, tenant_two,
+                dynamic_two[index].module, true)->clone(),
+        });
+    }
+
+    const int64_t launches_before_nan =
+        qwen36_get_dynamic_adam_launch_count(distributed);
+    if (rank == 0) {
+        setenv("QWEN36_TEST_INJECT_DYNAMIC_GRAD_NAN", "1", 1);
+    }
+    assert(qwen36_train_multi_lora_selected(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, &tenant_one, 1, kLoraRank) < 0.0);
+    unsetenv("QWEN36_TEST_INJECT_DYNAMIC_GRAD_NAN");
+    assert(qwen36_get_dynamic_adam_launch_count(distributed) ==
+        launches_before_nan);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) == 0);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_two) == 0);
+    for (size_t index = 0; index < dynamic_one.size(); ++index) {
+        const char* module = dynamic_one[index].module;
+        assert(max_diff(
+            *dynamic_tensor(distributed, tenant_one, module, false),
+            tenant_one_before[index][0]) == 0.0);
+        assert(max_diff(
+            *dynamic_tensor(distributed, tenant_one, module, true),
+            tenant_one_before[index][1]) == 0.0);
+        assert(dynamic_state(
+            distributed, tenant_one, module, false, false)
+            ->abs().max().item<double>() == 0.0);
+        assert(dynamic_state(
+            distributed, tenant_one, module, true, false)
+            ->abs().max().item<double>() == 0.0);
+    }
+
+    const double dynamic_loss = qwen36_train_multi_lora_selected(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, &tenant_one, 1, kLoraRank);
+    const double reference_dynamic_loss = qwen36_train_step(
+        dynamic_reference, &global_batch.input_ids, &global_batch.target_mask,
+        &global_batch.attention_mask);
+    assert(dynamic_loss > 0.0 && std::isfinite(dynamic_loss));
+    assert(reference_dynamic_loss > 0.0 && std::isfinite(reference_dynamic_loss));
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) == 1);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_two) == 0);
+    assert(qwen36_get_step_count(distributed) == fixed_step_before_dynamic);
+
+    double selected_update = 0.0;
+    double isolated_diff = 0.0;
+    double dynamic_adam_error = 0.0;
+    double dynamic_param_diff = 0.0;
+    double dynamic_m_diff = 0.0;
+    double dynamic_v_diff = 0.0;
+    std::vector<void*> dynamic_reference_m(kOptimizerSlots);
+    std::vector<void*> dynamic_reference_v(kOptimizerSlots);
+    assert(qwen36_export_optimizer_state(
+        dynamic_reference, dynamic_reference_m.data(),
+        dynamic_reference_v.data(), kOptimizerSlots) == kOptimizerSlots);
+    for (size_t index = 0; index < dynamic_one.size(); ++index) {
+        const char* module = dynamic_one[index].module;
+        auto* selected_a = dynamic_tensor(
+            distributed, tenant_one, module, false);
+        auto* selected_b = dynamic_tensor(
+            distributed, tenant_one, module, true);
+        auto* isolated_a = dynamic_tensor(
+            distributed, tenant_two, module, false);
+        auto* isolated_b = dynamic_tensor(
+            distributed, tenant_two, module, true);
+        const int64_t a_state = 2 * dynamic_one[index].slot;
+        const int64_t b_state = a_state + 1;
+        auto* reference_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(dynamic_reference, dynamic_one[index].slot));
+        auto* reference_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(dynamic_reference, dynamic_one[index].slot));
+        selected_update += update_norm(
+            *selected_a, tenant_one_before[index][0]);
+        selected_update += update_norm(
+            *selected_b, tenant_one_before[index][1]);
+        isolated_diff = std::max({isolated_diff,
+            max_diff(*isolated_a, tenant_two_before[index][0]),
+            max_diff(*isolated_b, tenant_two_before[index][1])});
+
+        auto* m_a = dynamic_state(
+            distributed, tenant_one, module, false, false);
+        auto* v_a = dynamic_state(
+            distributed, tenant_one, module, false, true);
+        auto* m_b = dynamic_state(
+            distributed, tenant_one, module, true, false);
+        auto* v_b = dynamic_state(
+            distributed, tenant_one, module, true, true);
+        auto* reference_m_a = reinterpret_cast<at::Tensor*>(
+            dynamic_reference_m[a_state]);
+        auto* reference_v_a = reinterpret_cast<at::Tensor*>(
+            dynamic_reference_v[a_state]);
+        auto* reference_m_b = reinterpret_cast<at::Tensor*>(
+            dynamic_reference_m[b_state]);
+        auto* reference_v_b = reinterpret_cast<at::Tensor*>(
+            dynamic_reference_v[b_state]);
+        assert(reference_a && reference_b && reference_m_a && reference_v_a &&
+            reference_m_b && reference_v_b);
+        dynamic_param_diff = std::max({dynamic_param_diff,
+            max_diff(*selected_a, local_factor(
+                *reference_a, dynamic_one[index].kind, false,
+                tp_rank, ep_rank, base_tp_mlp)),
+            max_diff(*selected_b, local_factor(
+                *reference_b, dynamic_one[index].kind, true,
+                tp_rank, ep_rank, base_tp_mlp))});
+        dynamic_m_diff = std::max({dynamic_m_diff,
+            max_diff(*m_a, local_factor(
+                *reference_m_a, dynamic_one[index].kind, false,
+                tp_rank, ep_rank, base_tp_mlp)),
+            max_diff(*m_b, local_factor(
+                *reference_m_b, dynamic_one[index].kind, true,
+                tp_rank, ep_rank, base_tp_mlp))});
+        dynamic_v_diff = std::max({dynamic_v_diff,
+            max_diff(*v_a, local_factor(
+                *reference_v_a, dynamic_one[index].kind, false,
+                tp_rank, ep_rank, base_tp_mlp)),
+            max_diff(*v_b, local_factor(
+                *reference_v_b, dynamic_one[index].kind, true,
+                tp_rank, ep_rank, base_tp_mlp))});
+        dynamic_adam_error = std::max({dynamic_adam_error,
+            max_diff(*selected_a, adam_expected(
+                tenant_one_before[index][0], *m_a, *v_a)),
+            max_diff(*selected_b, adam_expected(
+                tenant_one_before[index][1], *m_b, *v_b))});
+        assert(dynamic_state(
+            distributed, tenant_two, module, false, false)
+            ->abs().max().item<double>() == 0.0);
+        assert(dynamic_state(
+            distributed, tenant_two, module, true, false)
+            ->abs().max().item<double>() == 0.0);
+    }
+    std::printf(
+        "native_tp_ep_dynamic rank=%d tp=%d ep=%d dp=%d loss=%0.8f "
+        "selected_update=%0.8e isolated_diff=%0.8e adam_error=%0.8e "
+        "param_diff=%0.8e m_diff=%0.8e v_diff=%0.8e "
+        "steps=[%ld,%ld]\n",
+        rank, tp_rank, ep_rank, dp_rank, dynamic_loss, selected_update,
+        isolated_diff, dynamic_adam_error, dynamic_param_diff,
+        dynamic_m_diff, dynamic_v_diff,
+        static_cast<long>(qwen36_get_adapter_step_count(
+            distributed, tenant_one)),
+        static_cast<long>(qwen36_get_adapter_step_count(
+            distributed, tenant_two)));
+    std::fflush(stdout);
+    assert(selected_update > 0.0);
+    assert(isolated_diff == 0.0);
+    assert(dynamic_adam_error <= 1e-5);
+    assert(dynamic_param_diff <= 3e-3);
+    assert(dynamic_m_diff <= 7e-3);
+    assert(dynamic_v_diff <= 2e-5);
+
+    // Round-trip one tenant through CPU tensors and its independent optimizer
+    // clock, then require exact next-step parity with the uninterrupted context.
+    const auto first_host_snapshot = dynamic_state_cpu(
+        distributed, tenant_one, dynamic_one.front().module, false, false);
+    const auto repeated_host_snapshot = dynamic_state_cpu(
+        distributed, tenant_one, dynamic_one.front().module, false, false);
+    assert(max_diff(first_host_snapshot, repeated_host_snapshot) == 0.0);
+    std::vector<std::array<at::Tensor, 6>> dynamic_checkpoint;
+    dynamic_checkpoint.reserve(dynamic_one.size());
+    for (const auto& fixture : dynamic_one) {
+        dynamic_checkpoint.push_back({
+            dynamic_tensor(distributed, tenant_one, fixture.module, false)
+                ->to(at::kCPU).clone(),
+            dynamic_tensor(distributed, tenant_one, fixture.module, true)
+                ->to(at::kCPU).clone(),
+            dynamic_state_cpu(
+                distributed, tenant_one, fixture.module, false, false),
+            dynamic_state_cpu(
+                distributed, tenant_one, fixture.module, false, true),
+            dynamic_state_cpu(
+                distributed, tenant_one, fixture.module, true, false),
+            dynamic_state_cpu(
+                distributed, tenant_one, fixture.module, true, true),
+        });
+    }
+    const int64_t resumed_tenant_one = qwen36_add_lora_for_restore(
+        resumed, kLoraRank, kLoraRank, dynamic_targets, 1, targets);
+    const int64_t resumed_tenant_two = qwen36_add_lora_for_restore(
+        resumed, kLoraRank, kLoraRank, dynamic_targets, 1, targets);
+    assert(resumed_tenant_one == tenant_one);
+    assert(resumed_tenant_two == tenant_two);
+    std::vector<int64_t> checkpoint_layers;
+    std::vector<const char*> checkpoint_modules;
+    std::vector<void*> checkpoint_state_ptrs;
+    checkpoint_layers.reserve(dynamic_one.size());
+    checkpoint_modules.reserve(dynamic_one.size());
+    checkpoint_state_ptrs.reserve(dynamic_one.size() * 4);
+    for (size_t index = 0; index < dynamic_one.size(); ++index) {
+        const char* module = dynamic_one[index].module;
+        auto& state = dynamic_checkpoint[index];
+        assert(qwen36_set_adapter_lora_tensor(
+            resumed, resumed_tenant_one, 0, module, 0, &state[0]) == 0);
+        assert(qwen36_set_adapter_lora_tensor(
+            resumed, resumed_tenant_one, 0, module, 1, &state[1]) == 0);
+        checkpoint_layers.push_back(0);
+        checkpoint_modules.push_back(module);
+        checkpoint_state_ptrs.insert(checkpoint_state_ptrs.end(), {
+            &state[2], &state[3], &state[4], &state[5]});
+    }
+    // A partial descriptor set must fail before publishing a host snapshot or
+    // advancing the tenant clock.
+    assert(qwen36_import_adapter_optimizer_state_host_v1(
+        resumed, resumed_tenant_two, checkpoint_layers.data(),
+        checkpoint_modules.data(), checkpoint_state_ptrs.data(), 1, 1) < 0);
+    assert(qwen36_get_adapter_step_count(resumed, resumed_tenant_two) == 0);
+    assert(qwen36_get_adapter_optimizer_resident_count_v1(
+        resumed, resumed_tenant_two) == 0);
+
+    assert(qwen36_import_adapter_optimizer_state_host_v1(
+        resumed, resumed_tenant_one,
+        checkpoint_layers.data(), checkpoint_modules.data(),
+        checkpoint_state_ptrs.data(), checkpoint_layers.size(),
+        qwen36_get_adapter_step_count(distributed, tenant_one)) == 0);
+    assert(qwen36_get_adapter_step_count(resumed, resumed_tenant_one) ==
+        qwen36_get_adapter_step_count(distributed, tenant_one));
+    assert(qwen36_get_adapter_optimizer_resident_count_v1(
+        resumed, resumed_tenant_one) == 0);
+    double restored_dynamic_diff = 0.0;
+    for (const auto& fixture : dynamic_one) {
+        for (int is_b = 0; is_b < 2; ++is_b) {
+            restored_dynamic_diff = std::max(restored_dynamic_diff, max_diff(
+                *dynamic_tensor(distributed, tenant_one, fixture.module, is_b),
+                *dynamic_tensor(resumed, resumed_tenant_one, fixture.module, is_b)));
+            for (int is_v = 0; is_v < 2; ++is_v) {
+                const auto restored_state = dynamic_state_cpu(
+                    resumed, resumed_tenant_one, fixture.module, is_b, is_v);
+                restored_dynamic_diff = std::max(restored_dynamic_diff, max_diff(
+                    dynamic_state_cpu(
+                        distributed, tenant_one, fixture.module, is_b, is_v),
+                    restored_state));
+            }
+        }
+    }
+    assert(restored_dynamic_diff == 0.0);
+    assert(qwen36_get_adapter_optimizer_resident_count_v1(
+        resumed, resumed_tenant_one) == 0);
+    const int64_t fixed_step_before_resumed_dynamic =
+        qwen36_get_step_count(distributed);
+    const double continued_dynamic_loss = qwen36_train_multi_lora_selected(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, &tenant_one, 1, kLoraRank);
+    const double resumed_dynamic_loss = qwen36_train_multi_lora_selected(
+        resumed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, &resumed_tenant_one, 1, kLoraRank);
+    assert(std::abs(continued_dynamic_loss - resumed_dynamic_loss) <= 1e-6);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) == 2);
+    assert(qwen36_get_adapter_step_count(resumed, resumed_tenant_one) == 2);
+    assert(qwen36_get_adapter_step_count(resumed, resumed_tenant_two) == 0);
+    assert(qwen36_get_adapter_optimizer_resident_count_v1(
+        resumed, resumed_tenant_one) ==
+        static_cast<int64_t>(dynamic_one.size() * 4));
+    assert(qwen36_get_step_count(distributed) == fixed_step_before_resumed_dynamic);
+    assert(qwen36_get_step_count(resumed) == fixed_step_before_resumed_dynamic);
+    double resumed_dynamic_diff = 0.0;
+    for (const auto& fixture : dynamic_one) {
+        for (int is_b = 0; is_b < 2; ++is_b) {
+            resumed_dynamic_diff = std::max(resumed_dynamic_diff, max_diff(
+                *dynamic_tensor(distributed, tenant_one, fixture.module, is_b),
+                *dynamic_tensor(resumed, resumed_tenant_one, fixture.module, is_b)));
+            for (int is_v = 0; is_v < 2; ++is_v) {
+                resumed_dynamic_diff = std::max(resumed_dynamic_diff, max_diff(
+                    *dynamic_state(distributed, tenant_one, fixture.module, is_b, is_v),
+                    *dynamic_state(resumed, resumed_tenant_one, fixture.module, is_b, is_v)));
+            }
+        }
+    }
+    assert(resumed_dynamic_diff == 0.0);
+
+    // A zero-byte paging budget protects the selected tenant and evicts the
+    // colder resident tenant. The evicted tenant must hydrate from its host
+    // snapshot and preserve exact next-step parity with the resident oracle.
+    assert(qwen36_get_adapter_optimizer_resident_bytes_v1(
+        distributed, tenant_one) > 0);
+    assert(qwen36_get_adapter_optimizer_resident_bytes_v1(
+        distributed, tenant_two) > 0);
+    setenv("QWEN36_DYNAMIC_ADAM_HOST_PAGING", "1", 1);
+    setenv("QWEN36_DYNAMIC_ADAM_RESIDENT_BYTES", "0", 1);
+    const double paging_trigger_loss = qwen36_train_multi_lora_selected(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, &tenant_two, 1, kLoraRank);
+    assert(std::isfinite(paging_trigger_loss) && paging_trigger_loss >= 0.0);
+    assert(qwen36_get_adapter_optimizer_resident_bytes_v1(
+        distributed, tenant_one) == 0);
+    assert(qwen36_get_adapter_optimizer_resident_bytes_v1(
+        distributed, tenant_two) > 0);
+    const auto paged_tenant_one_m = dynamic_state_cpu(
+        distributed, tenant_one, dynamic_one.front().module, false, false);
+    assert(qwen36_get_adapter_optimizer_resident_bytes_v1(
+        distributed, tenant_one) == 0);
+    assert(max_diff(paged_tenant_one_m, dynamic_state_cpu(
+        resumed, resumed_tenant_one, dynamic_one.front().module,
+        false, false)) == 0.0);
+    unsetenv("QWEN36_DYNAMIC_ADAM_HOST_PAGING");
+    unsetenv("QWEN36_DYNAMIC_ADAM_RESIDENT_BYTES");
+
+    const double paged_next_loss = qwen36_train_multi_lora_selected(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, &tenant_one, 1, kLoraRank);
+    const double resident_next_loss = qwen36_train_multi_lora_selected(
+        resumed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, &resumed_tenant_one, 1, kLoraRank);
+    assert(std::abs(paged_next_loss - resident_next_loss) <= 1e-6);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) == 3);
+    assert(qwen36_get_adapter_step_count(resumed, resumed_tenant_one) == 3);
+    double paging_resume_diff = 0.0;
+    for (const auto& fixture : dynamic_one) {
+        for (int is_b = 0; is_b < 2; ++is_b) {
+            paging_resume_diff = std::max(paging_resume_diff, max_diff(
+                *dynamic_tensor(
+                    distributed, tenant_one, fixture.module, is_b),
+                *dynamic_tensor(
+                    resumed, resumed_tenant_one, fixture.module, is_b)));
+            for (int is_v = 0; is_v < 2; ++is_v) {
+                paging_resume_diff = std::max(paging_resume_diff, max_diff(
+                    *dynamic_state(distributed, tenant_one,
+                        fixture.module, is_b, is_v),
+                    *dynamic_state(resumed, resumed_tenant_one,
+                        fixture.module, is_b, is_v)));
+            }
+        }
+    }
+    assert(paging_resume_diff == 0.0);
+
+    const int64_t heterogeneous_tenant = qwen36_add_lora_v2(
+        distributed, 3, 3.0, &target_layer, 1, projection_targets);
+    assert(heterogeneous_tenant > tenant_two);
+    auto* heterogeneous_b = dynamic_tensor(
+        distributed, heterogeneous_tenant, projection_targets, true);
+    auto* homogeneous_b = dynamic_tensor(
+        distributed, tenant_one, projection_targets, true);
+    auto heterogeneous_before = heterogeneous_b->clone();
+    auto homogeneous_before = homogeneous_b->clone();
+    const int64_t heterogeneous_ids[] = {tenant_one, heterogeneous_tenant};
+    const int64_t tenant_one_step_before =
+        qwen36_get_adapter_step_count(distributed, tenant_one);
+
+    // TP peers are replicas of the same source row. A rank-local target mask
+    // change must fail on every topology coordinate before gradient sync.
+    auto mismatched_tp_targets = local_batch.target_mask.clone();
+    if (ep_rank == 0 &&
+        ((dp_rank == 0 && tp_rank == 0) ||
+         (dp_rank == 1 && tp_rank == 1))) {
+        mismatched_tp_targets.zero_();
+    }
+    const int64_t finalizers_before_tp_mismatch =
+        qwen36_get_dynamic_finalizer_count(distributed);
+    const int64_t adam_before_tp_mismatch =
+        qwen36_get_dynamic_adam_launch_count(distributed);
+    assert(qwen36_train_multi_lora_selected_v2(
+        distributed, &local_batch.input_ids, &mismatched_tp_targets,
+        &local_batch.attention_mask, heterogeneous_ids, 2) < 0.0);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) ==
+        tenant_one_step_before);
+    assert(qwen36_get_adapter_step_count(distributed, heterogeneous_tenant) == 0);
+    assert(max_diff(*homogeneous_b, homogeneous_before) == 0.0);
+    assert(max_diff(*heterogeneous_b, heterogeneous_before) == 0.0);
+    assert(qwen36_get_dynamic_finalizer_count(distributed) ==
+        finalizers_before_tp_mismatch + 1);
+    assert(qwen36_get_dynamic_adam_launch_count(distributed) ==
+        adam_before_tp_mismatch);
+
+    const int64_t finalizers_before_heterogeneous =
+        qwen36_get_dynamic_finalizer_count(distributed);
+    const int64_t adam_launches_before_heterogeneous =
+        qwen36_get_dynamic_adam_launch_count(distributed);
+
+    if (rank == 0)
+        setenv("QWEN36_TEST_FAIL_HETERO_GROUP_AFTER", "1", 1);
+    assert(qwen36_train_multi_lora_selected_v2(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, heterogeneous_ids, 2) < 0.0);
+    if (rank == 0)
+        unsetenv("QWEN36_TEST_FAIL_HETERO_GROUP_AFTER");
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) ==
+        tenant_one_step_before);
+    assert(qwen36_get_adapter_step_count(distributed, heterogeneous_tenant) == 0);
+    assert(max_diff(*homogeneous_b, homogeneous_before) == 0.0);
+    assert(max_diff(*heterogeneous_b, heterogeneous_before) == 0.0);
+    assert(qwen36_get_dynamic_finalizer_count(distributed) ==
+        finalizers_before_heterogeneous);
+    assert(qwen36_get_dynamic_adam_launch_count(distributed) ==
+        adam_launches_before_heterogeneous);
+
+    if (rank == 0)
+        setenv("QWEN36_TEST_FAIL_FINALIZER_BEFORE_TOKEN_PREFLIGHT", "1", 1);
+    assert(qwen36_train_multi_lora_selected_v2(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, heterogeneous_ids, 2) < 0.0);
+    if (rank == 0)
+        unsetenv("QWEN36_TEST_FAIL_FINALIZER_BEFORE_TOKEN_PREFLIGHT");
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) ==
+        tenant_one_step_before);
+    assert(qwen36_get_adapter_step_count(distributed, heterogeneous_tenant) == 0);
+    assert(max_diff(*homogeneous_b, homogeneous_before) == 0.0);
+    assert(max_diff(*heterogeneous_b, heterogeneous_before) == 0.0);
+    assert(qwen36_get_dynamic_finalizer_count(distributed) ==
+        finalizers_before_heterogeneous + 1);
+    assert(qwen36_get_dynamic_adam_launch_count(distributed) ==
+        adam_launches_before_heterogeneous);
+
+    if (rank == 0)
+        setenv("QWEN36_TEST_FAIL_DYNAMIC_ADAM_BEFORE_COMMIT", "1", 1);
+    assert(qwen36_train_multi_lora_selected_v2(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, heterogeneous_ids, 2) < 0.0);
+    if (rank == 0)
+        unsetenv("QWEN36_TEST_FAIL_DYNAMIC_ADAM_BEFORE_COMMIT");
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) ==
+        tenant_one_step_before);
+    assert(qwen36_get_adapter_step_count(distributed, heterogeneous_tenant) == 0);
+    assert(max_diff(*homogeneous_b, homogeneous_before) == 0.0);
+    assert(max_diff(*heterogeneous_b, heterogeneous_before) == 0.0);
+    assert(qwen36_get_dynamic_finalizer_count(distributed) ==
+        finalizers_before_heterogeneous + 2);
+    assert(qwen36_get_dynamic_adam_launch_count(distributed) ==
+        adam_launches_before_heterogeneous + 1);
+
+    // Exercise request-order token normalization across two signatures with
+    // unequal positive counts and a third globally-empty tenant. The empty
+    // tenant must retain its parameters, Adam state, and private clock.
+    const int64_t heterogeneous_batch_ids[] = {
+        tenant_one, heterogeneous_tenant, tenant_two};
+    auto dense_targets = at::ones_like(local_batch.target_mask);
+    dense_targets.select(1, 0).zero_();
+    auto empty_targets = at::zeros_like(local_batch.target_mask);
+    auto heterogeneous_input = local_batch.input_ids.repeat({3, 1});
+    auto heterogeneous_targets = at::cat({
+        local_batch.target_mask, dense_targets, empty_targets}, 0);
+    auto heterogeneous_attention =
+        local_batch.attention_mask.repeat({3, 1});
+    auto* empty_tenant_b = dynamic_tensor(
+        distributed, tenant_two, projection_targets, true);
+    auto* empty_tenant_m = dynamic_state(
+        distributed, tenant_two, projection_targets, true, false);
+    auto* empty_tenant_v = dynamic_state(
+        distributed, tenant_two, projection_targets, true, true);
+    assert(empty_tenant_b && empty_tenant_m && empty_tenant_v);
+    const auto empty_tenant_b_before = empty_tenant_b->clone();
+    const auto empty_tenant_m_before = empty_tenant_m->clone();
+    const auto empty_tenant_v_before = empty_tenant_v->clone();
+    const int64_t empty_tenant_step_before =
+        qwen36_get_adapter_step_count(distributed, tenant_two);
+    const int64_t heterogeneous_train_batches_before =
+        qwen36_get_dynamic_train_batch_count(distributed);
+
+    double mixed_aggregate = -1.0;
+    double mixed_adapter_losses[3] = {-1.0, -1.0, -1.0};
+    if (rank == 0) {
+        assert(qwen36_train_multi_lora_selected_v3(
+            distributed, &heterogeneous_input, &heterogeneous_targets,
+            &heterogeneous_attention, heterogeneous_batch_ids, 3,
+            &mixed_aggregate, mixed_adapter_losses, 3) < 0);
+        assert(mixed_aggregate == -1.0);
+        for (const double adapter_loss : mixed_adapter_losses)
+            assert(adapter_loss == -1.0);
+    } else {
+        assert(qwen36_train_multi_lora_selected_v2(
+            distributed, &heterogeneous_input, &heterogeneous_targets,
+            &heterogeneous_attention, heterogeneous_batch_ids, 3) < 0.0);
+    }
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) ==
+        tenant_one_step_before);
+    assert(qwen36_get_adapter_step_count(distributed, heterogeneous_tenant) == 0);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_two) ==
+        empty_tenant_step_before);
+
+    double heterogeneous_loss = -1.0;
+    double heterogeneous_adapter_losses[3] = {-1.0, -1.0, -1.0};
+    assert(qwen36_train_multi_lora_selected_v3(
+        distributed, &heterogeneous_input, &heterogeneous_targets,
+        &heterogeneous_attention, heterogeneous_batch_ids, 3,
+        &heterogeneous_loss, heterogeneous_adapter_losses, 3) == 0);
+    assert(heterogeneous_loss > 0.0 && std::isfinite(heterogeneous_loss));
+    for (const double adapter_loss : heterogeneous_adapter_losses)
+        assert(adapter_loss >= 0.0 && std::isfinite(adapter_loss));
+    const double reported_mean =
+        (heterogeneous_adapter_losses[0] +
+         heterogeneous_adapter_losses[1] +
+         heterogeneous_adapter_losses[2]) / 3.0;
+    assert(std::abs(heterogeneous_loss - reported_mean) < 1e-8);
+    assert(heterogeneous_adapter_losses[2] == 0.0);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_one) ==
+        tenant_one_step_before + 1);
+    assert(qwen36_get_adapter_step_count(distributed, heterogeneous_tenant) == 1);
+    assert(qwen36_get_adapter_step_count(distributed, tenant_two) ==
+        empty_tenant_step_before);
+    assert(qwen36_get_dynamic_finalizer_count(distributed) ==
+        finalizers_before_heterogeneous + 3);
+    assert(qwen36_get_dynamic_adam_launch_count(distributed) ==
+        adam_launches_before_heterogeneous + 2);
+    const char* heterogeneous_padded_env =
+        std::getenv("QWEN36_HETERO_PADDED_BATCH");
+    const int64_t expected_heterogeneous_train_batches =
+        heterogeneous_padded_env &&
+            std::strcmp(heterogeneous_padded_env, "0") == 0
+        ? 2
+        : 1;
+    assert(qwen36_get_dynamic_train_batch_count(distributed) ==
+        heterogeneous_train_batches_before +
+            expected_heterogeneous_train_batches);
+    assert(update_norm(*homogeneous_b, homogeneous_before) > 0.0);
+    assert(update_norm(*heterogeneous_b, heterogeneous_before) > 0.0);
+    assert(max_diff(*empty_tenant_b, empty_tenant_b_before) == 0.0);
+    assert(max_diff(*empty_tenant_m, empty_tenant_m_before) == 0.0);
+    assert(max_diff(*empty_tenant_v, empty_tenant_v_before) == 0.0);
+    std::printf(
+        "native_tp_ep_heterogeneous_v2 rank=%d loss=%0.8f ok\n",
+        rank, heterogeneous_loss);
+    std::fflush(stdout);
+
+    const int64_t poison_request[] = {tenant_one + 100000};
+    if (rank == 0)
+        setenv("QWEN36_TEST_POISON_DYNAMIC_RECOVERY", "1", 1);
+    assert(qwen36_train_multi_lora_selected_v2(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, poison_request, 1) < 0.0);
+    if (rank == 0)
+        unsetenv("QWEN36_TEST_POISON_DYNAMIC_RECOVERY");
+    assert(qwen36_train_multi_lora_selected_v2(
+        distributed, &local_batch.input_ids, &local_batch.target_mask,
+        &local_batch.attention_mask, &tenant_one, 1) < 0.0);
+    assert(qwen36_get_context_health(distributed) == -1);
+    std::printf("native_tp_ep_dynamic_recovery_poison rank=%d ok\n", rank);
+    std::fflush(stdout);
+
+    qwen36_free_training_context(reference);
+    qwen36_free_training_context(dynamic_reference);
+    qwen36_free_training_context(resumed);
+    qwen36_free_training_context(distributed);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp b/crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp
new file mode 100644
index 00000000..ecd93d57
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp
@@ -0,0 +1,406 @@
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDACachingAllocator.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <iomanip>
+#include <initializer_list>
+#include <numeric>
+#include <sstream>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" int64_t qwen36_kernel_abi_version();
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_nccl(void*);
+extern "C" int32_t qwen36_init_parallel_nccl_v2(
+    void*, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t,
+    int32_t, int32_t, int32_t);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" void qwen36_free_training_context(void*);
+
+namespace {
+
+constexpr int64_t kAbiVersion = 31;
+constexpr int32_t kBaseTpAttention = 1 << 0;
+constexpr int32_t kVocabParallel = 1 << 2;
+
+static int env_int(const char* name, int fallback) {
+    const char* value = std::getenv(name);
+    if (!value || value[0] == '\0') return fallback;
+    const int parsed = std::atoi(value);
+    assert(parsed > 0);
+    return parsed;
+}
+
+static int env_int_or(const char* primary, const char* secondary, int fallback) {
+    return std::getenv(primary) ? env_int(primary, fallback)
+                                : env_int(secondary, fallback);
+}
+
+static bool env_enabled(const char* name, bool fallback = false) {
+    const char* value = std::getenv(name);
+    if (!value || value[0] == '\0') return fallback;
+    return std::strcmp(value, "0") != 0 && std::strcmp(value, "false") != 0;
+}
+
+static at::Tensor seeded_randn(
+    std::initializer_list<int64_t> shape, double scale, int64_t seed
+) {
+    at::manual_seed(seed);
+    return (at::randn(shape,
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)) * scale)
+        .to(at::kBFloat16);
+}
+
+static at::Tensor unit_weight(std::initializer_list<int64_t> shape) {
+    return at::ones(
+        shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static void append_gdn_layer(
+    std::vector<at::Tensor>& weights, int64_t layer,
+    int64_t hidden, int64_t intermediate,
+    int64_t global_k_heads, int64_t global_v_heads,
+    int64_t key_dim, int64_t value_dim, int64_t conv_kernel,
+    int tp_world
+) {
+    const int64_t local_k_heads = global_k_heads / tp_world;
+    const int64_t local_v_heads = global_v_heads / tp_world;
+    const int64_t q_size = local_k_heads * key_dim;
+    const int64_t v_size = local_v_heads * value_dim;
+    const int64_t qkv_size = 2 * q_size + v_size;
+    const int64_t seed = 1000 + layer * 100;
+
+    weights.push_back(unit_weight({hidden}));
+    weights.push_back(unit_weight({hidden}));
+    weights.push_back(seeded_randn({qkv_size, hidden}, 0.0020, seed + 2));
+    weights.push_back(seeded_randn({v_size, hidden}, 0.0020, seed + 3));
+    weights.push_back(seeded_randn({local_v_heads, hidden}, 0.0020, seed + 4));
+    weights.push_back(seeded_randn({local_v_heads, hidden}, 0.0020, seed + 5));
+    // A realistic negative time-step bias keeps the recurrent decay near one.
+    // A zero bias would produce decay ~= 0.5 and make reverse state recovery
+    // exponentially ill-conditioned over this synthetic long sequence.
+    weights.push_back(at::zeros({local_v_heads},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    weights.push_back(at::full({local_v_heads}, -4.0,
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    weights.push_back(seeded_randn(
+        {qkv_size, 1, conv_kernel}, 0.0020, seed + 8));
+    weights.push_back(unit_weight({value_dim}));
+    weights.push_back(seeded_randn({hidden, v_size}, 0.0020, seed + 10));
+    weights.push_back(seeded_randn(
+        {intermediate, hidden}, 0.0020, seed + 11));
+    weights.push_back(seeded_randn(
+        {intermediate, hidden}, 0.0020, seed + 12));
+    weights.push_back(seeded_randn(
+        {hidden, intermediate}, 0.0020, seed + 13));
+}
+
+static double percentile(std::vector<double> values, double quantile) {
+    assert(!values.empty());
+    std::sort(values.begin(), values.end());
+    const double position = quantile * static_cast<double>(values.size() - 1);
+    const size_t lower = static_cast<size_t>(position);
+    const size_t upper = std::min(lower + 1, values.size() - 1);
+    const double fraction = position - static_cast<double>(lower);
+    return values[lower] * (1.0 - fraction) + values[upper] * fraction;
+}
+
+static double gib(size_t bytes) {
+    return static_cast<double>(bytes) / (1024.0 * 1024.0 * 1024.0);
+}
+
+static size_t used_since(size_t initial_free, size_t observed_free) {
+    return initial_free > observed_free ? initial_free - observed_free : 0;
+}
+
+}  // namespace
+
+int main() {
+    const char* mode_env = std::getenv("BENCH_MODE");
+    const std::string mode = mode_env ? mode_env : "single";
+    assert(mode == "single" || mode == "tp2" || mode == "cp2");
+    const bool use_tp = mode == "tp2";
+    const bool use_cp = mode == "cp2";
+    const int expected_world = (use_tp || use_cp) ? 2 : 1;
+    const int weight_shards = use_tp ? expected_world : 1;
+    const int rank = std::getenv("RANK") ? std::atoi(std::getenv("RANK")) : 0;
+    const int world = std::getenv("WORLD_SIZE")
+        ? std::atoi(std::getenv("WORLD_SIZE")) : 1;
+    const int local_rank = std::getenv("LOCAL_RANK")
+        ? std::atoi(std::getenv("LOCAL_RANK")) : rank;
+    assert(world == expected_world && rank >= 0 && rank < world);
+    assert(qwen36_kernel_abi_version() == kAbiVersion);
+    qwen36_set_cuda_device(local_rank);
+    assert(cudaFree(nullptr) == cudaSuccess);
+
+    const int batch = env_int("BENCH_BATCH", 2);
+    const int seq = env_int("BENCH_SEQ", 512);
+    const int hidden = env_int("BENCH_HIDDEN", 2048);
+    const int key_heads = env_int("BENCH_K_HEADS", 16);
+    const int value_heads = env_int("BENCH_V_HEADS", 32);
+    const int key_dim = env_int_or("BENCH_KEY_DIM", "BENCH_HEAD_DIM", 128);
+    const int value_dim = env_int_or("BENCH_VALUE_DIM", "BENCH_HEAD_DIM", 128);
+    const int conv_kernel = env_int("BENCH_CONV", 4);
+    const int layers = env_int("BENCH_LAYERS", 3);
+    const int intermediate = env_int("BENCH_INTERMEDIATE", 2048);
+    const int vocab = env_int("BENCH_VOCAB", 4096);
+    const int lora_rank = env_int("BENCH_LORA_RANK", 8);
+    const int warmup = env_int("BENCH_WARMUP", 5);
+    const int iters = env_int("BENCH_ITERS", 30);
+    const bool fused_ab = env_enabled("QWEN36_GDN_FUSED_AB_PROJECTION", true);
+    const bool fused_cp_exchange = env_enabled(
+        "QWEN36_GDN_FUSED_CP_EXCHANGE");
+    const bool recurrent_fusion = env_enabled(
+        "QWEN36_GDN_RECURRENT_FUSION", true);
+    assert(seq >= 2 && hidden > 0 && key_dim > 0 && value_dim > 0);
+    assert(value_heads % key_heads == 0);
+    assert(key_heads % expected_world == 0 && value_heads % expected_world == 0);
+    assert(!use_tp || lora_rank % expected_world == 0);
+    assert(!use_tp || vocab % expected_world == 0);
+
+    c10::cuda::CUDACachingAllocator::resetPeakStats(local_rank);
+    size_t free_start = 0, total_bytes = 0;
+    assert(cudaMemGetInfo(&free_start, &total_bytes) == cudaSuccess);
+    size_t min_observed_free = free_start;
+
+    std::vector<at::Tensor> weights;
+    weights.reserve(layers * 14);
+    for (int64_t layer = 0; layer < layers; ++layer) {
+        append_gdn_layer(weights, layer, hidden, intermediate,
+            key_heads, value_heads, key_dim, value_dim, conv_kernel,
+            weight_shards);
+    }
+    for (auto& weight : weights) weight.set_requires_grad(false);
+    auto weight_ptrs = pointers(weights);
+
+    const int local_vocab = vocab / weight_shards;
+    const int rank_seed_offset = use_tp ? rank : 0;
+    auto embed = seeded_randn(
+        {local_vocab, hidden}, 0.0020, 31 + rank_seed_offset);
+    auto final_norm = unit_weight({hidden});
+    auto lm_head = seeded_randn(
+        {local_vocab, hidden}, 0.0020, 37 + rank_seed_offset);
+    embed.set_requires_grad(false);
+    final_norm.set_requires_grad(false);
+    lm_head.set_requires_grad(false);
+
+    std::vector<LayerConfig> configs(layers);
+    for (auto& config : configs) {
+        config.layer_type = 1;
+        config.num_k_heads = key_heads;
+        config.key_dim = key_dim;
+        config.num_v_heads = value_heads;
+        config.val_dim = value_dim;
+        config.conv_kernel = conv_kernel;
+        config.rms_eps = 1e-5;
+        config.intermediate_size = intermediate;
+    }
+    std::vector<int64_t> target_layers(layers);
+    std::iota(target_layers.begin(), target_layers.end(), 0);
+    constexpr const char* targets =
+        "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj";
+
+    std::vector<int64_t> host_ids(batch * seq);
+    for (int b = 0; b < batch; ++b) {
+        for (int s = 0; s < seq; ++s)
+            host_ids[b * seq + s] = 1 + (b * 131 + s * 17) % (vocab - 1);
+    }
+    auto input_ids = at::from_blob(host_ids.data(), {batch, seq},
+        at::TensorOptions().device(at::kCPU).dtype(at::kLong)).clone().to(at::kCUDA);
+    auto target_mask = at::ones({batch, seq},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    auto attention_mask = at::ones({batch, seq},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+
+    size_t free_before_context = 0;
+    assert(cudaMemGetInfo(&free_before_context, &total_bytes) == cudaSuccess);
+    min_observed_free = std::min(min_observed_free, free_before_context);
+    setenv("TP_SIZE", use_tp ? "2" : "1", 1);
+    setenv("CP_SIZE", use_cp ? "2" : "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("PP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", use_tp ? std::to_string(rank).c_str() : "0", 1);
+    setenv("RUSTRAIN_CP_RANK", use_cp ? std::to_string(rank).c_str() : "0", 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    setenv("RUSTRAIN_PP_RANK", "0", 1);
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    void* context = (use_tp || use_cp)
+        ? qwen36_create_training_context_ex(
+            weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head,
+            configs.data(), layers, static_cast<int32_t>(at::kBFloat16),
+            1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+            target_layers.data(), layers, targets,
+            use_tp ? kBaseTpAttention | kVocabParallel : 0)
+        : qwen36_create_training_context(
+            weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head,
+            configs.data(), layers, static_cast<int32_t>(at::kBFloat16),
+            1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+            target_layers.data(), layers, targets);
+    assert(context);
+    if (use_tp) assert(qwen36_init_nccl(context) == 0);
+    if (use_cp) {
+        assert(qwen36_init_parallel_nccl_v2(
+            context, rank, world,
+            0, 1, 0,
+            rank, 2, 0,
+            0, 1, 0,
+            0, 1, 0,
+            0, 1, 0) == 0);
+    }
+
+    size_t free_after_context = 0;
+    assert(cudaMemGetInfo(&free_after_context, &total_bytes) == cudaSuccess);
+    min_observed_free = std::min(min_observed_free, free_after_context);
+    double last_loss = 0.0;
+    for (int i = 0; i < warmup; ++i) {
+        last_loss = qwen36_train_step(
+            context, &input_ids, &target_mask, &attention_mask);
+        assert(std::isfinite(last_loss));
+    }
+    assert(cudaDeviceSynchronize() == cudaSuccess);
+    c10::cuda::CUDACachingAllocator::resetPeakStats(local_rank);
+
+    size_t free_after_warmup = 0;
+    assert(cudaMemGetInfo(&free_after_warmup, &total_bytes) == cudaSuccess);
+    min_observed_free = std::min(min_observed_free, free_after_warmup);
+
+    cudaEvent_t start = nullptr;
+    cudaEvent_t stop = nullptr;
+    assert(cudaEventCreate(&start) == cudaSuccess);
+    assert(cudaEventCreate(&stop) == cudaSuccess);
+    std::vector<double> times;
+    times.reserve(iters);
+    auto stream = c10::cuda::getCurrentCUDAStream().stream();
+    for (int i = 0; i < iters; ++i) {
+        assert(cudaEventRecord(start, stream) == cudaSuccess);
+        last_loss = qwen36_train_step(
+            context, &input_ids, &target_mask, &attention_mask);
+        assert(cudaEventRecord(stop, stream) == cudaSuccess);
+        assert(cudaEventSynchronize(stop) == cudaSuccess);
+        assert(std::isfinite(last_loss));
+        float elapsed_ms = 0.0f;
+        assert(cudaEventElapsedTime(&elapsed_ms, start, stop) == cudaSuccess);
+        times.push_back(elapsed_ms);
+        size_t free_now = 0;
+        assert(cudaMemGetInfo(&free_now, &total_bytes) == cudaSuccess);
+        min_observed_free = std::min(min_observed_free, free_now);
+    }
+    assert(cudaEventDestroy(start) == cudaSuccess);
+    assert(cudaEventDestroy(stop) == cudaSuccess);
+
+    const double mean = std::accumulate(times.begin(), times.end(), 0.0) /
+        static_cast<double>(times.size());
+    double variance = 0.0;
+    for (const double value : times) variance += (value - mean) * (value - mean);
+    variance /= static_cast<double>(times.size());
+    const double p50 = percentile(times, 0.50);
+    const double p90 = percentile(times, 0.90);
+    const double model_tokens = static_cast<double>(batch) * (seq - 1);
+    const double layer_tokens = model_tokens * layers;
+    const double model_tokens_per_second = model_tokens / (p50 / 1000.0);
+    const double layer_tokens_per_second = layer_tokens / (p50 / 1000.0);
+
+    const auto allocator_stats =
+        c10::cuda::CUDACachingAllocator::getDeviceStats(local_rank);
+    constexpr size_t aggregate = 0;
+    const auto& allocated = allocator_stats.allocated_bytes[aggregate];
+    const auto& reserved = allocator_stats.reserved_bytes[aggregate];
+
+    cudaDeviceProp properties{};
+    assert(cudaGetDeviceProperties(&properties, local_rank) == cudaSuccess);
+    std::ostringstream output;
+    output << std::fixed << std::setprecision(6)
+        << "native_tp_gdn_bench {"
+        << "\"mode\":\"" << mode << "\","
+        << "\"rank\":" << rank << ",\"world\":" << world << ","
+        << "\"gpu\":\"" << properties.name << "\","
+        << "\"abi\":" << kAbiVersion << ","
+        << "\"fused_ab_projection\":" << (fused_ab ? "true" : "false") << ","
+        << "\"fused_cp_exchange\":"
+        << (fused_cp_exchange ? "true" : "false") << ","
+        << "\"gdn_recurrent_fusion\":"
+        << (recurrent_fusion ? "true" : "false") << ","
+        << "\"batch\":" << batch << ",\"seq\":" << seq << ","
+        << "\"hidden\":" << hidden << ",\"layers\":" << layers << ","
+        << "\"key_heads\":" << key_heads << ","
+        << "\"value_heads\":" << value_heads << ","
+        << "\"key_dim\":" << key_dim << ","
+        << "\"value_dim\":" << value_dim << ","
+        << "\"n_rep\":" << (value_heads / key_heads) << ","
+        << "\"conv\":" << conv_kernel << ","
+        << "\"intermediate\":" << intermediate << ","
+        << "\"vocab\":" << vocab << ",\"lora_rank\":" << lora_rank << ","
+        << "\"warmup\":" << warmup << ",\"iters\":" << iters << ","
+        << "\"last_loss\":" << last_loss << ","
+        << "\"step_ms_mean\":" << mean << ","
+        << "\"step_ms_p50\":" << p50 << ","
+        << "\"step_ms_p90\":" << p90 << ","
+        << "\"step_ms_std\":" << std::sqrt(variance) << ","
+        << "\"model_tokens\":" << model_tokens << ","
+        << "\"gdn_layer_tokens\":" << layer_tokens << ","
+        << "\"model_tokens_per_sec\":" << model_tokens_per_second << ","
+        << "\"gdn_layer_tokens_per_sec\":" << layer_tokens_per_second << ","
+        << "\"device_total_gib\":" << gib(total_bytes) << ","
+        << "\"free_start_gib\":" << gib(free_start) << ","
+        << "\"free_before_context_gib\":" << gib(free_before_context) << ","
+        << "\"free_after_context_gib\":" << gib(free_after_context) << ","
+        << "\"free_after_warmup_gib\":" << gib(free_after_warmup) << ","
+        << "\"context_resident_delta_gib\":"
+        << gib(used_since(free_before_context, free_after_context)) << ","
+        << "\"max_observed_resident_gib\":"
+        << gib(used_since(free_start, min_observed_free)) << ","
+        << "\"allocator_peak_allocated_gib\":" << gib(allocated.peak) << ','
+        << "\"allocator_peak_reserved_gib\":" << gib(reserved.peak) << ','
+        << "\"samples_ms\":[";
+    for (size_t i = 0; i < times.size(); ++i) {
+        if (i) output << ',';
+        output << times[i];
+    }
+    output << "]}\n";
+    const std::string line = output.str();
+    std::fwrite(line.data(), 1, line.size(), stdout);
+    std::fflush(stdout);
+
+    qwen36_free_training_context(context);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp
new file mode 100644
index 00000000..d56248d9
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp
@@ -0,0 +1,993 @@
+#include <ATen/ATen.h>
+
+#include <algorithm>
+#include <array>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <initializer_list>
+#include <string>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" int64_t qwen36_kernel_abi_version();
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_nccl(void*);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" int64_t qwen36_list_lora(void*, int64_t*, int64_t);
+extern "C" void* qwen36_get_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_set_adapter_lora_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, void*);
+extern "C" void* qwen36_get_adapter_optimizer_tensor(
+    void*, int64_t, int64_t, const char*, int32_t, int32_t);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t);
+extern "C" int32_t qwen36_abort_gradient_accumulation(void*);
+extern "C" int64_t qwen36_export_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" double qwen36_eval_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_micro_step(
+    void*, void*, void*, void*, double, int32_t);
+extern "C" int32_t qwen36_train_multi_lora_selected_v3(
+    void*, void*, void*, void*, const int64_t*, int32_t,
+    double*, double*, int32_t);
+extern "C" void qwen36_set_checkpoint(void*, int32_t, int64_t);
+extern "C" int64_t qwen36_get_lora_batch_projection_build_count(void*);
+extern "C" int64_t qwen36_get_lora_batch_scaling_upload_count(void*);
+extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t);
+extern "C" void qwen36_free_training_context(void*);
+
+namespace {
+
+constexpr int64_t kAbiVersion = 31;
+constexpr int32_t kBaseTpAttention = 1 << 0;
+constexpr int32_t kVocabParallel = 1 << 2;
+constexpr int64_t kLayers = 2;
+constexpr int64_t kHidden = 32;
+constexpr int64_t kIntermediate = 48;
+constexpr int64_t kKeyHeads = 4;
+constexpr int64_t kValueHeads = 8;
+constexpr int64_t kKeyDim = 128;
+constexpr int64_t kValueDim = 128;
+constexpr int64_t kConvKernel = 4;
+constexpr int64_t kVocab = 64;
+constexpr int64_t kLoraRank = 4;
+constexpr int64_t kPairsPerLayer = 8;
+constexpr int64_t kQSize = kKeyHeads * kKeyDim;
+constexpr int64_t kVSize = kValueHeads * kValueDim;
+constexpr int64_t kQkvSize = 2 * kQSize + kVSize;
+
+constexpr std::array<const char*, 5> kGdnModules = {
+    "in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b", "out_proj"};
+
+static int required_env_int(const char* name) {
+    const char* value = std::getenv(name);
+    assert(value && value[0] != '\0');
+    return std::atoi(value);
+}
+
+static at::Tensor fingerprint(
+    std::initializer_list<int64_t> shape, double scale, int64_t offset
+) {
+    int64_t count = 1;
+    for (const int64_t dim : shape) count *= dim;
+    auto values = at::arange(
+        count, at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    return ((values.add(offset).remainder(97) - 48.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static at::Tensor unit_weight(std::initializer_list<int64_t> shape) {
+    return at::ones(
+        shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+}
+
+static void append_full_gdn_layer(
+    std::vector<at::Tensor>& weights, int64_t layer
+) {
+    const int64_t base = 1000 * layer;
+    weights.push_back(unit_weight({kHidden}));
+    weights.push_back(unit_weight({kHidden}));
+
+    auto q = fingerprint({kQSize, kHidden}, 0.00035, base + 11);
+    auto k = fingerprint({kQSize, kHidden}, 0.00041, base + 211);
+    auto v = fingerprint({kVSize, kHidden}, 0.00029, base + 421);
+    weights.push_back(at::cat({q, k, v}, 0).contiguous());
+    weights.push_back(fingerprint({kVSize, kHidden}, 0.00031, base + 631));
+    weights.push_back(fingerprint({kValueHeads, kHidden}, 0.00043, base + 719));
+    weights.push_back(fingerprint({kValueHeads, kHidden}, 0.00047, base + 811));
+    weights.push_back(fingerprint({kValueHeads}, 0.0008, base + 907));
+    weights.push_back(fingerprint({kValueHeads}, 0.0007, base + 953));
+
+    auto q_conv = fingerprint(
+        {kQSize, 1, kConvKernel}, 0.0009, base + 101);
+    auto k_conv = fingerprint(
+        {kQSize, 1, kConvKernel}, 0.0011, base + 307);
+    auto v_conv = fingerprint(
+        {kVSize, 1, kConvKernel}, 0.0007, base + 509);
+    weights.push_back(at::cat({q_conv, k_conv, v_conv}, 0).contiguous());
+    weights.push_back(unit_weight({kValueDim}));
+    weights.push_back(fingerprint({kHidden, kVSize}, 0.00033, base + 613));
+
+    weights.push_back(fingerprint(
+        {kIntermediate, kHidden}, 0.00045, base + 701));
+    weights.push_back(fingerprint(
+        {kIntermediate, kHidden}, 0.00039, base + 797));
+    weights.push_back(fingerprint(
+        {kHidden, kIntermediate}, 0.00037, base + 887));
+}
+
+static at::Tensor shard_flat_qkv(const at::Tensor& full, int rank, int dim) {
+    const int64_t local_q = kQSize / 2;
+    const int64_t local_v = kVSize / 2;
+    auto q = full.narrow(dim, rank * local_q, local_q);
+    auto k = full.narrow(dim, kQSize + rank * local_q, local_q);
+    auto v = full.narrow(dim, 2 * kQSize + rank * local_v, local_v);
+    return at::cat({q, k, v}, dim).contiguous();
+}
+
+static std::vector<at::Tensor> make_local_weights(
+    const std::vector<at::Tensor>& full, int rank
+) {
+    std::vector<at::Tensor> local;
+    local.reserve(full.size());
+    for (int64_t layer = 0; layer < kLayers; ++layer) {
+        const int64_t offset = layer * 14;
+        local.push_back(full[offset + 0]);
+        local.push_back(full[offset + 1]);
+        local.push_back(shard_flat_qkv(full[offset + 2], rank, 0));
+        local.push_back(full[offset + 3]
+            .narrow(0, rank * (kVSize / 2), kVSize / 2).contiguous());
+        local.push_back(full[offset + 4]
+            .narrow(0, rank * (kValueHeads / 2), kValueHeads / 2).contiguous());
+        local.push_back(full[offset + 5]
+            .narrow(0, rank * (kValueHeads / 2), kValueHeads / 2).contiguous());
+        local.push_back(full[offset + 6]
+            .narrow(0, rank * (kValueHeads / 2), kValueHeads / 2).contiguous());
+        local.push_back(full[offset + 7]
+            .narrow(0, rank * (kValueHeads / 2), kValueHeads / 2).contiguous());
+        local.push_back(shard_flat_qkv(full[offset + 8], rank, 0));
+        local.push_back(full[offset + 9]);
+        local.push_back(full[offset + 10]
+            .narrow(1, rank * (kVSize / 2), kVSize / 2).contiguous());
+        local.insert(local.end(), full.begin() + offset + 11,
+            full.begin() + offset + 14);
+    }
+    return local;
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static LayerConfig gdn_config() {
+    LayerConfig config{};
+    config.layer_type = 1;
+    config.num_k_heads = kKeyHeads;
+    config.key_dim = kKeyDim;
+    config.num_v_heads = kValueHeads;
+    config.val_dim = kValueDim;
+    config.conv_kernel = kConvKernel;
+    config.rms_eps = 1e-5;
+    config.intermediate_size = kIntermediate;
+    return config;
+}
+
+struct Batch {
+    at::Tensor input_ids;
+    at::Tensor target_mask;
+    at::Tensor attention_mask;
+};
+
+static Batch make_batch(int64_t batch, int64_t seq, int64_t offset) {
+    std::vector<int64_t> host(batch * seq);
+    for (int64_t b = 0; b < batch; ++b) {
+        for (int64_t s = 0; s < seq; ++s) {
+            host[b * seq + s] = 1 + (offset + b * 17 + s * 5) % (kVocab - 1);
+        }
+    }
+    auto ids = at::from_blob(host.data(), {batch, seq},
+        at::TensorOptions().device(at::kCPU).dtype(at::kLong)).clone().to(at::kCUDA);
+    auto target = at::ones({batch, seq},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    auto attention = at::ones({batch, seq},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+    return {std::move(ids), std::move(target), std::move(attention)};
+}
+
+static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) {
+    assert(lhs.sizes() == rhs.sizes());
+    return (lhs.to(at::kFloat) - rhs.to(at::kFloat))
+        .abs().max().item<double>();
+}
+
+static double relative_l2(const at::Tensor& lhs, const at::Tensor& rhs) {
+    assert(lhs.sizes() == rhs.sizes());
+    auto delta = lhs.to(at::kFloat) - rhs.to(at::kFloat);
+    const double denom = std::max(rhs.to(at::kFloat).norm().item<double>(), 1e-12);
+    return delta.norm().item<double>() / denom;
+}
+
+static bool is_column_parallel(const char* module) {
+    return std::strcmp(module, "out_proj") != 0;
+}
+
+static at::Tensor expected_first_adam_step(
+    const at::Tensor& parameter,
+    const at::Tensor& first_moment,
+    const at::Tensor& second_moment,
+    double effective_beta1,
+    double effective_beta2
+) {
+    constexpr double lr = 1e-3;
+    constexpr double eps = 1e-8;
+    const float lr_scaled = static_cast<float>(
+        lr * std::sqrt(1.0 - effective_beta2) /
+        (1.0 - effective_beta1));
+    const float eps_scaled = static_cast<float>(
+        eps * std::sqrt(1.0 - effective_beta2));
+    return (parameter.to(at::kFloat) -
+        lr_scaled * first_moment /
+            (second_moment.sqrt() + eps_scaled)).to(at::kBFloat16);
+}
+
+static int64_t full_output_size(const char* module) {
+    if (std::strcmp(module, "in_proj_qkv") == 0) return kQkvSize;
+    if (std::strcmp(module, "in_proj_z") == 0) return kVSize;
+    if (std::strcmp(module, "in_proj_a") == 0 ||
+        std::strcmp(module, "in_proj_b") == 0) return kValueHeads;
+    assert(std::strcmp(module, "out_proj") == 0);
+    return kHidden;
+}
+
+static int64_t full_input_size(const char* module) {
+    return std::strcmp(module, "out_proj") == 0 ? kVSize : kHidden;
+}
+
+static at::Tensor shard_projection_output(
+    const at::Tensor& full, const char* module, int rank, int dim
+) {
+    if (std::strcmp(module, "in_proj_qkv") == 0)
+        return shard_flat_qkv(full, rank, dim);
+    const int64_t local = full.size(dim) / 2;
+    return full.narrow(dim, rank * local, local).contiguous();
+}
+
+static at::Tensor expected_local_factor(
+    const at::Tensor& full, const char* module, bool is_b, int rank
+) {
+    if (is_column_parallel(module)) {
+        return is_b ? shard_projection_output(full, module, rank, 0)
+                    : full;
+    }
+    return is_b ? full
+                : full.narrow(1, rank * (kVSize / 2), kVSize / 2).contiguous();
+}
+
+struct LoraFixture {
+    int64_t layer;
+    int64_t pair;
+    const char* module;
+    at::Tensor full_a;
+    at::Tensor full_b;
+    at::Tensor local_a;
+    at::Tensor local_b;
+};
+
+static std::vector<LoraFixture> make_lora_fixtures(int rank, int64_t seed_base) {
+    std::vector<LoraFixture> fixtures;
+    fixtures.reserve(kLayers * kGdnModules.size());
+    for (int64_t layer = 0; layer < kLayers; ++layer) {
+        for (int64_t pair = 0; pair < static_cast<int64_t>(kGdnModules.size()); ++pair) {
+            const char* module = kGdnModules[pair];
+            const int64_t base = seed_base + layer * 100 + pair * 13;
+            auto full_a = fingerprint(
+                {kLoraRank, full_input_size(module)}, 0.0007, base + 1);
+            auto full_b = fingerprint(
+                {full_output_size(module), kLoraRank}, 0.0006, base + 7);
+            auto local_a = expected_local_factor(full_a, module, false, rank);
+            auto local_b = expected_local_factor(full_b, module, true, rank);
+            fixtures.push_back({layer, pair, module, std::move(full_a),
+                std::move(full_b), std::move(local_a), std::move(local_b)});
+        }
+    }
+    return fixtures;
+}
+
+static void assert_weight_contract(
+    const std::vector<at::Tensor>& full,
+    const std::vector<at::Tensor>& local, int rank
+) {
+    assert(full.size() == 28 && local.size() == full.size());
+    for (int64_t layer = 0; layer < kLayers; ++layer) {
+        const int64_t offset = layer * 14;
+        assert(local[offset + 2].sizes() ==
+            at::IntArrayRef({kQkvSize / 2, kHidden}));
+        assert(local[offset + 8].sizes() ==
+            at::IntArrayRef({kQkvSize / 2, 1, kConvKernel}));
+        assert(local[offset + 3].sizes() ==
+            at::IntArrayRef({kVSize / 2, kHidden}));
+        assert(local[offset + 4].sizes() ==
+            at::IntArrayRef({kValueHeads / 2, kHidden}));
+        assert(local[offset + 9].sizes() == at::IntArrayRef({kValueDim}));
+        assert(local[offset + 10].sizes() ==
+            at::IntArrayRef({kHidden, kVSize / 2}));
+        assert(max_diff(local[offset + 2],
+            shard_flat_qkv(full[offset + 2], rank, 0)) == 0.0);
+        assert(max_diff(local[offset + 8],
+            shard_flat_qkv(full[offset + 8], rank, 0)) == 0.0);
+    }
+}
+
+struct ContextPair {
+    void* distributed = nullptr;
+    void* reference = nullptr;
+};
+
+static ContextPair create_context_pair(
+    std::vector<at::Tensor>& local_weights,
+    std::vector<at::Tensor>& full_weights,
+    at::Tensor& local_embed, at::Tensor& full_embed,
+    at::Tensor& final_norm,
+    at::Tensor& local_lm_head, at::Tensor& full_lm_head,
+    LayerConfig* configs, const char* fixed_targets,
+    bool run_negative_guard
+) {
+    const int64_t target_layers[kLayers] = {0, 1};
+    auto local_ptrs = pointers(local_weights);
+    setenv("TP_SIZE", "2", 1);
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+
+    if (run_negative_guard) {
+        auto invalid_weights = local_weights;
+        invalid_weights[2] = invalid_weights[2]
+            .narrow(0, 0, invalid_weights[2].size(0) - 1).contiguous();
+        auto invalid_ptrs = pointers(invalid_weights);
+        void* invalid = qwen36_create_training_context_ex(
+            invalid_ptrs.data(), invalid_ptrs.size(), &local_embed, &final_norm,
+            &local_lm_head,
+            configs, kLayers, static_cast<int32_t>(at::kBFloat16),
+            1.0, 1e-3, 0.9, 0.999, 1e-8, kVocab, 1e-5, kLoraRank,
+            target_layers, kLayers, fixed_targets,
+            kBaseTpAttention | kVocabParallel);
+        assert(!invalid && "invalid local flat-QKV shape must be rejected");
+
+        invalid = qwen36_create_training_context_ex(
+            local_ptrs.data(), local_ptrs.size(), &full_embed, &final_norm,
+            &local_lm_head,
+            configs, kLayers, static_cast<int32_t>(at::kBFloat16),
+            1.0, 1e-3, 0.9, 0.999, 1e-8, kVocab, 1e-5, kLoraRank,
+            target_layers, kLayers, fixed_targets,
+            kBaseTpAttention | kVocabParallel);
+        assert(!invalid && "invalid local embedding shape must be rejected");
+    }
+
+    void* distributed = qwen36_create_training_context_ex(
+        local_ptrs.data(), local_ptrs.size(), &local_embed, &final_norm,
+        &local_lm_head,
+        configs, kLayers, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, kVocab, 1e-5, kLoraRank,
+        target_layers, kLayers, fixed_targets,
+        kBaseTpAttention | kVocabParallel);
+    assert(distributed && qwen36_init_nccl(distributed) == 0);
+
+    const std::string distributed_rank = std::getenv("RANK");
+    const std::string distributed_world = std::getenv("WORLD_SIZE");
+    setenv("WORLD_SIZE", "1", 1);
+    setenv("RANK", "0", 1);
+    setenv("TP_SIZE", "1", 1);
+    setenv("EP_SIZE", "1", 1);
+    setenv("DP_SIZE", "1", 1);
+    setenv("RUSTRAIN_TP_RANK", "0", 1);
+    setenv("RUSTRAIN_EP_RANK", "0", 1);
+    setenv("RUSTRAIN_DP_RANK", "0", 1);
+    auto full_ptrs = pointers(full_weights);
+    void* reference = qwen36_create_training_context(
+        full_ptrs.data(), full_ptrs.size(), &full_embed, &final_norm,
+        &full_lm_head,
+        configs, kLayers, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, kVocab, 1e-5, kLoraRank,
+        target_layers, kLayers, fixed_targets);
+    assert(reference);
+    setenv("WORLD_SIZE", distributed_world.c_str(), 1);
+    setenv("RANK", distributed_rank.c_str(), 1);
+    setenv("TP_SIZE", "2", 1);
+    unsetenv("RUSTRAIN_TP_RANK");
+    unsetenv("RUSTRAIN_EP_RANK");
+    unsetenv("RUSTRAIN_DP_RANK");
+    return {distributed, reference};
+}
+
+static void set_fixed_lora(
+    ContextPair contexts, std::vector<LoraFixture>& fixtures
+) {
+    for (auto& fixture : fixtures) {
+        const int64_t slot = fixture.layer * kPairsPerLayer + fixture.pair;
+        assert(qwen36_set_lora_tensor(
+            contexts.distributed, slot, 0, &fixture.local_a) == 0);
+        assert(qwen36_set_lora_tensor(
+            contexts.distributed, slot, 1, &fixture.local_b) == 0);
+        assert(qwen36_set_lora_tensor(
+            contexts.reference, slot, 0, &fixture.full_a) == 0);
+        assert(qwen36_set_lora_tensor(
+            contexts.reference, slot, 1, &fixture.full_b) == 0);
+    }
+}
+
+struct ErrorSummary {
+    double a = 0.0;
+    double b = 0.0;
+    double relative = 0.0;
+    double m = 0.0;
+    double v = 0.0;
+    double param = 0.0;
+    double adam = 0.0;
+};
+
+static void check_backward_reference_parity(
+    void* context, const std::vector<LoraFixture>& fixtures,
+    std::vector<at::Tensor>& local_weights, Batch& batch, int rank
+) {
+    struct SavedEnv {
+        const char* name;
+        bool present;
+        std::string value;
+    };
+    std::vector<SavedEnv> saved_env;
+    for (const char* name : {
+             "QWEN36_DELTA_REFERENCE_BWD", "QWEN36_GDN_INVERSE_BWD",
+             "QWEN36_GDN_CHUNKWISE_BWD",
+             "QWEN36_GDN_STATE_CHECKPOINT_STRIDE"}) {
+        const char* value = std::getenv(name);
+        saved_env.push_back({name, value != nullptr, value ? value : ""});
+    }
+    unsetenv("QWEN36_DELTA_REFERENCE_BWD");
+    unsetenv("QWEN36_GDN_INVERSE_BWD");
+    unsetenv("QWEN36_GDN_CHUNKWISE_BWD");
+    setenv("QWEN36_GDN_STATE_CHECKPOINT_STRIDE", "4", 1);
+
+    // Stress both ordinary and near-zero recurrence decays without changing
+    // the baseline TP fixture used by the rest of this smoke test.
+    auto stress_a_log = at::tensor(
+        {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 4.75f},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat))
+        .to(at::kBFloat16);
+    std::vector<at::Tensor> original_a_log;
+    original_a_log.reserve(kLayers);
+    for (int64_t layer = 0; layer < kLayers; ++layer) {
+        auto& local_a_log = local_weights[layer * 14 + 6];
+        original_a_log.push_back(local_a_log.clone());
+        local_a_log.copy_(stress_a_log.narrow(
+            0, rank * local_a_log.size(0), local_a_log.size(0)));
+    }
+
+    const double stable_loss = qwen36_train_micro_step(
+        context, &batch.input_ids, &batch.target_mask,
+        &batch.attention_mask, 1.0, 0);
+    assert(std::isfinite(stable_loss));
+    std::vector<at::Tensor> stable_a;
+    std::vector<at::Tensor> stable_b;
+    stable_a.reserve(fixtures.size());
+    stable_b.reserve(fixtures.size());
+    for (const auto& fixture : fixtures) {
+        const int64_t slot = fixture.layer * kPairsPerLayer + fixture.pair;
+        auto* a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_grad_accumulator(context, slot, 0));
+        auto* b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_grad_accumulator(context, slot, 1));
+        assert(a && b && a->scalar_type() == at::kFloat &&
+            b->scalar_type() == at::kFloat);
+        stable_a.push_back(a->clone());
+        stable_b.push_back(b->clone());
+    }
+    assert(qwen36_abort_gradient_accumulation(context) == 0);
+
+    setenv("QWEN36_DELTA_REFERENCE_BWD", "1", 1);
+    const double reference_loss = qwen36_train_micro_step(
+        context, &batch.input_ids, &batch.target_mask,
+        &batch.attention_mask, 1.0, 0);
+    assert(std::isfinite(reference_loss));
+
+    double worst_absolute = 0.0;
+    double worst_relative = 0.0;
+    double worst_reference = 0.0;
+    int64_t mismatches = 0;
+    int64_t significant_sign_mismatches = 0;
+    int64_t compared = 0;
+    for (size_t index = 0; index < fixtures.size(); ++index) {
+        const auto& fixture = fixtures[index];
+        const int64_t slot = fixture.layer * kPairsPerLayer + fixture.pair;
+        for (int is_b = 0; is_b < 2; ++is_b) {
+            auto* reference = reinterpret_cast<at::Tensor*>(
+                qwen36_get_lora_grad_accumulator(context, slot, is_b));
+            assert(reference && reference->scalar_type() == at::kFloat);
+            const auto& stable = is_b ? stable_b[index] : stable_a[index];
+            auto stable_f = stable.to(at::kFloat);
+            auto reference_f = reference->to(at::kFloat);
+            auto delta = (stable_f - reference_f).abs();
+            worst_reference = std::max(
+                worst_reference, reference_f.abs().max().item<double>());
+            worst_absolute = std::max(
+                worst_absolute, delta.max().item<double>());
+            worst_relative = std::max(
+                worst_relative, relative_l2(stable_f, reference_f));
+            auto close = at::isclose(
+                stable_f, reference_f, 1e-2, 1e-4, true);
+            mismatches += close.logical_not().sum().item<int64_t>();
+            auto sign_mismatch = stable_f.sign().ne(reference_f.sign());
+            significant_sign_mismatches += sign_mismatch
+                .logical_and(reference_f.abs().gt(1e-5))
+                .sum().item<int64_t>();
+            compared += reference_f.numel();
+        }
+    }
+    assert(qwen36_abort_gradient_accumulation(context) == 0);
+    for (int64_t layer = 0; layer < kLayers; ++layer) {
+        local_weights[layer * 14 + 6].copy_(original_a_log[layer]);
+    }
+    for (const auto& saved : saved_env) {
+        if (saved.present) {
+            setenv(saved.name, saved.value.c_str(), 1);
+        } else {
+            unsetenv(saved.name);
+        }
+    }
+
+    std::printf(
+        "native_gdn_backward_parity rank=%d loss_diff=%0.8e "
+        "max_abs=%0.8e max_ref=%0.8e max_relative_l2=%0.8e "
+        "mismatches=%ld/%ld "
+        "significant_sign_mismatches=%ld\n",
+        rank, std::abs(stable_loss - reference_loss), worst_absolute,
+        worst_reference, worst_relative, mismatches, compared,
+        significant_sign_mismatches);
+    std::fflush(stdout);
+    assert(std::abs(stable_loss - reference_loss) < 1e-8);
+    assert(mismatches == 0);
+    assert(significant_sign_mismatches == 0);
+}
+
+static void check_fixed_path(
+    ContextPair contexts, std::vector<LoraFixture>& fixtures,
+    Batch& short_batch, Batch& batch, int rank
+) {
+    const double short_distributed = qwen36_eval_step(contexts.distributed,
+        &short_batch.input_ids, &short_batch.target_mask,
+        &short_batch.attention_mask);
+    const double short_reference = qwen36_eval_step(contexts.reference,
+        &short_batch.input_ids, &short_batch.target_mask,
+        &short_batch.attention_mask);
+    const double eval_distributed = qwen36_eval_step(contexts.distributed,
+        &batch.input_ids, &batch.target_mask, &batch.attention_mask);
+    const double eval_reference = qwen36_eval_step(contexts.reference,
+        &batch.input_ids, &batch.target_mask, &batch.attention_mask);
+    assert(std::isfinite(short_distributed) && std::isfinite(short_reference));
+    assert(std::isfinite(eval_distributed) && std::isfinite(eval_reference));
+
+    const double micro_distributed = qwen36_train_micro_step(
+        contexts.distributed, &batch.input_ids, &batch.target_mask,
+        &batch.attention_mask, 1.0, 0);
+    const double micro_reference = qwen36_train_micro_step(
+        contexts.reference, &batch.input_ids, &batch.target_mask,
+        &batch.attention_mask, 1.0, 0);
+    assert(std::isfinite(micro_distributed) && std::isfinite(micro_reference));
+
+    ErrorSummary errors;
+    for (auto& fixture : fixtures) {
+        const int64_t slot = fixture.layer * kPairsPerLayer + fixture.pair;
+        auto* local_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_grad_accumulator(contexts.distributed, slot, 0));
+        auto* local_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_grad_accumulator(contexts.distributed, slot, 1));
+        auto* full_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_grad_accumulator(contexts.reference, slot, 0));
+        auto* full_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_grad_accumulator(contexts.reference, slot, 1));
+        assert(local_a && local_b && full_a && full_b);
+        assert(local_a->scalar_type() == at::kFloat &&
+            local_b->scalar_type() == at::kFloat);
+        auto expected_a = expected_local_factor(*full_a, fixture.module, false, rank);
+        auto expected_b = expected_local_factor(*full_b, fixture.module, true, rank);
+        const double a_diff = max_diff(*local_a, expected_a);
+        const double b_diff = max_diff(*local_b, expected_b);
+        const double a_relative = relative_l2(*local_a, expected_a);
+        const double b_relative = relative_l2(*local_b, expected_b);
+        std::printf(
+            "native_tp_gdn_grad rank=%d layer=%ld module=%s "
+            "a_diff=%0.8e b_diff=%0.8e a_relative_l2=%0.8e "
+            "b_relative_l2=%0.8e a_reference_norm=%0.8e "
+            "b_reference_norm=%0.8e\n",
+            rank, fixture.layer, fixture.module, a_diff, b_diff,
+            a_relative, b_relative,
+            expected_a.to(at::kFloat).norm().item<double>(),
+            expected_b.to(at::kFloat).norm().item<double>());
+        // Before the final optimizer boundary, replicated factors have only
+        // their rank-local contribution. Compare shard-local factors here;
+        // final Adam state checks below cover the synchronized replicas.
+        if (is_column_parallel(fixture.module)) {
+            errors.b = std::max(errors.b, b_diff);
+            errors.relative = std::max(errors.relative, b_relative);
+        } else {
+            errors.a = std::max(errors.a, a_diff);
+            errors.relative = std::max(errors.relative, a_relative);
+        }
+    }
+    assert(qwen36_abort_gradient_accumulation(contexts.distributed) == 0);
+    assert(qwen36_abort_gradient_accumulation(contexts.reference) == 0);
+
+    const double loss_distributed = qwen36_train_step(contexts.distributed,
+        &batch.input_ids, &batch.target_mask, &batch.attention_mask);
+    const double loss_reference = qwen36_train_step(contexts.reference,
+        &batch.input_ids, &batch.target_mask, &batch.attention_mask);
+    assert(std::isfinite(loss_distributed) && std::isfinite(loss_reference));
+
+    constexpr int64_t optimizer_count = 2 * kLayers * kPairsPerLayer;
+    std::vector<void*> local_m(optimizer_count), local_v(optimizer_count);
+    std::vector<void*> full_m(optimizer_count), full_v(optimizer_count);
+    assert(qwen36_export_optimizer_state(contexts.distributed,
+        local_m.data(), local_v.data(), optimizer_count) == optimizer_count);
+    assert(qwen36_export_optimizer_state(contexts.reference,
+        full_m.data(), full_v.data(), optimizer_count) == optimizer_count);
+
+    for (auto& fixture : fixtures) {
+        const int64_t slot = fixture.layer * kPairsPerLayer + fixture.pair;
+        auto* updated_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(contexts.distributed, slot));
+        auto* updated_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(contexts.distributed, slot));
+        auto* reference_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(contexts.reference, slot));
+        auto* reference_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(contexts.reference, slot));
+        assert(updated_a && updated_b && reference_a && reference_b);
+
+        auto* m_a = reinterpret_cast<at::Tensor*>(local_m[2 * slot]);
+        auto* m_b = reinterpret_cast<at::Tensor*>(local_m[2 * slot + 1]);
+        auto* v_a = reinterpret_cast<at::Tensor*>(local_v[2 * slot]);
+        auto* v_b = reinterpret_cast<at::Tensor*>(local_v[2 * slot + 1]);
+        auto* ref_m_a = reinterpret_cast<at::Tensor*>(full_m[2 * slot]);
+        auto* ref_m_b = reinterpret_cast<at::Tensor*>(full_m[2 * slot + 1]);
+        auto* ref_v_a = reinterpret_cast<at::Tensor*>(full_v[2 * slot]);
+        auto* ref_v_b = reinterpret_cast<at::Tensor*>(full_v[2 * slot + 1]);
+        assert(m_a && m_b && v_a && v_b && ref_m_a && ref_m_b &&
+            ref_v_a && ref_v_b);
+
+        errors.m = std::max({errors.m,
+            max_diff(*m_a, expected_local_factor(
+                *ref_m_a, fixture.module, false, rank)),
+            max_diff(*m_b, expected_local_factor(
+                *ref_m_b, fixture.module, true, rank))});
+        errors.v = std::max({errors.v,
+            max_diff(*v_a, expected_local_factor(
+                *ref_v_a, fixture.module, false, rank)),
+            max_diff(*v_b, expected_local_factor(
+                *ref_v_b, fixture.module, true, rank))});
+        errors.param = std::max({errors.param,
+            max_diff(*updated_a, expected_local_factor(
+                *reference_a, fixture.module, false, rank)),
+            max_diff(*updated_b, expected_local_factor(
+                *reference_b, fixture.module, true, rank))});
+
+        auto expected_a = expected_first_adam_step(
+            fixture.local_a, *m_a, *v_a, 0.9, 0.999);
+        auto expected_b = expected_first_adam_step(
+            fixture.local_b, *m_b, *v_b, 0.9, 0.999);
+        errors.adam = std::max({errors.adam,
+            max_diff(*updated_a, expected_a), max_diff(*updated_b, expected_b)});
+    }
+
+    const double short_diff = std::abs(short_distributed - short_reference);
+    const double eval_diff = std::abs(eval_distributed - eval_reference);
+    const double micro_diff = std::abs(micro_distributed - micro_reference);
+    const double loss_diff = std::abs(loss_distributed - loss_reference);
+    std::printf(
+        "native_tp_gdn_fixed rank=%d short_eval_diff=%0.8e eval_diff=%0.8e "
+        "micro_diff=%0.8e loss_diff=%0.8e a_grad_diff=%0.8e "
+        "b_grad_diff=%0.8e grad_relative_l2=%0.8e m_diff=%0.8e "
+        "v_diff=%0.8e param_diff=%0.8e adam_error=%0.8e\n",
+        rank, short_diff, eval_diff, micro_diff, loss_diff, errors.a, errors.b,
+        errors.relative, errors.m, errors.v, errors.param, errors.adam);
+    std::fflush(stdout);
+
+    assert(short_diff < 5e-3 && eval_diff < 5e-3);
+    assert(micro_diff < 5e-3 && loss_diff < 5e-3);
+    assert(errors.a < 5e-4 && errors.b < 5e-4);
+    // Canonical Q/K normalization adds epsilon inside the squared norm. For
+    // the tiny in_proj_b fixture this exposes one extra BF16 TP rounding level
+    // in relative error; the absolute gradient bound above remains unchanged.
+    assert(errors.relative < 3.1e-2);
+    assert(errors.m < 5e-5 && errors.v < 5e-8);
+    // The row-local BF16 head dgrad can land one quantization level away from
+    // the full-vocabulary reference while the FP32 optimizer oracle stays tight.
+    assert(errors.param <= 2.1e-3);
+    assert(errors.adam < 1e-8);
+}
+
+static void set_dynamic_lora(
+    void* context, int64_t adapter_id,
+    std::vector<LoraFixture>& fixtures, bool local
+) {
+    for (auto& fixture : fixtures) {
+        auto& a = local ? fixture.local_a : fixture.full_a;
+        auto& b = local ? fixture.local_b : fixture.full_b;
+        assert(qwen36_set_adapter_lora_tensor(context, adapter_id,
+            fixture.layer, fixture.module, 0, &a) == 0);
+        assert(qwen36_set_adapter_lora_tensor(context, adapter_id,
+            fixture.layer, fixture.module, 1, &b) == 0);
+    }
+}
+
+static void check_dynamic_path(
+    ContextPair contexts, std::vector<LoraFixture>& fixtures,
+    Batch& batch, int rank
+) {
+    const int64_t target_layers[kLayers] = {0, 1};
+    constexpr const char* targets =
+        "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj";
+    const int64_t distributed_id = qwen36_add_lora(contexts.distributed,
+        kLoraRank, kLoraRank, target_layers, kLayers, targets);
+    const int64_t reference_id = qwen36_add_lora(contexts.reference,
+        kLoraRank, kLoraRank, target_layers, kLayers, targets);
+    assert(distributed_id > 0 && reference_id > 0);
+    set_dynamic_lora(contexts.distributed, distributed_id, fixtures, true);
+    set_dynamic_lora(contexts.reference, reference_id, fixtures, false);
+
+    const char* chunkwise_value = std::getenv("QWEN36_GDN_CHUNKWISE_BWD");
+    if (chunkwise_value && chunkwise_value[0] != '\0' &&
+        std::strcmp(chunkwise_value, "0") != 0 &&
+        std::strcmp(chunkwise_value, "false") != 0) {
+        const std::string saved_chunkwise_value(chunkwise_value);
+        if (rank == 0) unsetenv("QWEN36_GDN_CHUNKWISE_BWD");
+        double mismatched_loss = -1.0;
+        double mismatched_adapter_loss = -1.0;
+        assert(qwen36_train_multi_lora_selected_v3(
+            contexts.distributed, &batch.input_ids, &batch.target_mask,
+            &batch.attention_mask, &distributed_id, 1, &mismatched_loss,
+            &mismatched_adapter_loss, 1) < 0);
+        assert(mismatched_loss == -1.0 && mismatched_adapter_loss == -1.0);
+        assert(qwen36_get_adapter_step_count(
+            contexts.distributed, distributed_id) == 0);
+        if (rank == 0) {
+            setenv("QWEN36_GDN_CHUNKWISE_BWD",
+                saved_chunkwise_value.c_str(), 1);
+        }
+    }
+
+    double distributed_loss = -1.0;
+    double distributed_adapter_loss = -1.0;
+    double reference_loss = -1.0;
+    double reference_adapter_loss = -1.0;
+    assert(qwen36_train_multi_lora_selected_v3(
+        contexts.distributed, &batch.input_ids, &batch.target_mask,
+        &batch.attention_mask, &distributed_id, 1, &distributed_loss,
+        &distributed_adapter_loss, 1) == 0);
+    assert(qwen36_train_multi_lora_selected_v3(
+        contexts.reference, &batch.input_ids, &batch.target_mask,
+        &batch.attention_mask, &reference_id, 1, &reference_loss,
+        &reference_adapter_loss, 1) == 0);
+    assert(std::isfinite(distributed_loss) && std::isfinite(reference_loss));
+    assert(std::abs(distributed_loss - distributed_adapter_loss) < 1e-8);
+    assert(std::abs(reference_loss - reference_adapter_loss) < 1e-8);
+
+    ErrorSummary errors;
+    for (auto& fixture : fixtures) {
+        auto* local_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_lora_tensor(contexts.distributed,
+                distributed_id, fixture.layer, fixture.module, 0));
+        auto* local_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_lora_tensor(contexts.distributed,
+                distributed_id, fixture.layer, fixture.module, 1));
+        auto* full_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_lora_tensor(contexts.reference,
+                reference_id, fixture.layer, fixture.module, 0));
+        auto* full_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_adapter_lora_tensor(contexts.reference,
+                reference_id, fixture.layer, fixture.module, 1));
+        assert(local_a && local_b && full_a && full_b);
+        errors.param = std::max({errors.param,
+            max_diff(*local_a, expected_local_factor(
+                *full_a, fixture.module, false, rank)),
+            max_diff(*local_b, expected_local_factor(
+                *full_b, fixture.module, true, rank))});
+
+        auto state = [&](void* context, int64_t adapter, bool is_b, bool is_v) {
+            auto* tensor = reinterpret_cast<at::Tensor*>(
+                qwen36_get_adapter_optimizer_tensor(context, adapter,
+                    fixture.layer, fixture.module, is_b, is_v));
+            assert(tensor && tensor->scalar_type() == at::kFloat);
+            return tensor;
+        };
+        auto* m_a = state(contexts.distributed, distributed_id, false, false);
+        auto* m_b = state(contexts.distributed, distributed_id, true, false);
+        auto* v_a = state(contexts.distributed, distributed_id, false, true);
+        auto* v_b = state(contexts.distributed, distributed_id, true, true);
+        auto* full_m_a = state(contexts.reference, reference_id, false, false);
+        auto* full_m_b = state(contexts.reference, reference_id, true, false);
+        auto* full_v_a = state(contexts.reference, reference_id, false, true);
+        auto* full_v_b = state(contexts.reference, reference_id, true, true);
+        errors.m = std::max({errors.m,
+            max_diff(*m_a, expected_local_factor(
+                *full_m_a, fixture.module, false, rank)),
+            max_diff(*m_b, expected_local_factor(
+                *full_m_b, fixture.module, true, rank))});
+        errors.v = std::max({errors.v,
+            max_diff(*v_a, expected_local_factor(
+                *full_v_a, fixture.module, false, rank)),
+            max_diff(*v_b, expected_local_factor(
+                *full_v_b, fixture.module, true, rank))});
+
+        auto expected_a = expected_first_adam_step(
+            fixture.local_a, *m_a, *v_a,
+            static_cast<double>(0.9f), static_cast<double>(0.999f));
+        auto expected_b = expected_first_adam_step(
+            fixture.local_b, *m_b, *v_b,
+            static_cast<double>(0.9f), static_cast<double>(0.999f));
+        errors.adam = std::max({errors.adam,
+            max_diff(*local_a, expected_a), max_diff(*local_b, expected_b)});
+    }
+
+    const double loss_diff = std::abs(distributed_loss - reference_loss);
+    std::printf(
+        "native_tp_gdn_dynamic rank=%d loss_diff=%0.8e m_diff=%0.8e "
+        "v_diff=%0.8e param_diff=%0.8e adam_error=%0.8e\n",
+        rank, loss_diff, errors.m, errors.v, errors.param, errors.adam);
+    std::fflush(stdout);
+    assert(loss_diff < 5e-3);
+    assert(errors.m < 5e-5 && errors.v < 5e-8);
+    // Distributed online softmax can move the BF16 Adam result by one extra
+    // quantization level even when the FP32 m/v oracles remain much tighter.
+    assert(errors.param <= 2.1e-3);
+    assert(errors.adam < 1e-8);
+}
+
+}  // namespace
+
+int main() {
+    const int rank = required_env_int("RANK");
+    const int world = required_env_int("WORLD_SIZE");
+    const int local_rank = std::getenv("LOCAL_RANK")
+        ? std::atoi(std::getenv("LOCAL_RANK")) : rank;
+    assert(world == 2 && rank >= 0 && rank < world);
+    assert(qwen36_kernel_abi_version() == kAbiVersion);
+    qwen36_set_cuda_device(local_rank);
+
+    std::vector<at::Tensor> full_weights;
+    full_weights.reserve(kLayers * 14);
+    for (int64_t layer = 0; layer < kLayers; ++layer)
+        append_full_gdn_layer(full_weights, layer);
+    for (auto& weight : full_weights) weight.set_requires_grad(false);
+    auto local_weights = make_local_weights(full_weights, rank);
+    for (auto& weight : local_weights) weight.set_requires_grad(false);
+    assert_weight_contract(full_weights, local_weights, rank);
+
+    auto embed = fingerprint({kVocab, kHidden}, 0.0013, 17);
+    auto final_norm = unit_weight({kHidden});
+    auto lm_head = fingerprint({kVocab, kHidden}, 0.0011, 59);
+    embed.set_requires_grad(false);
+    final_norm.set_requires_grad(false);
+    lm_head.set_requires_grad(false);
+    auto local_embed = embed.narrow(
+        0, rank * (kVocab / world), kVocab / world).contiguous();
+    auto local_lm_head = lm_head.narrow(
+        0, rank * (kVocab / world), kVocab / world).contiguous();
+    LayerConfig configs[kLayers] = {gdn_config(), gdn_config()};
+    auto short_batch = make_batch(2, 3, 3);
+    auto batch = make_batch(2, 9, 11);
+    auto dynamic_batch = make_batch(1, 9, 23);
+
+    constexpr const char* all_targets =
+        "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj";
+    auto fixed_contexts = create_context_pair(local_weights, full_weights,
+        local_embed, embed, final_norm, local_embed, embed,
+        configs, all_targets, true);
+    auto fixed_fixtures = make_lora_fixtures(rank, 2000);
+    set_fixed_lora(fixed_contexts, fixed_fixtures);
+    check_backward_reference_parity(
+        fixed_contexts.distributed, fixed_fixtures, local_weights, batch, rank);
+    // GDN state is recurrent, so left padding/internal holes must be rejected
+    // until the native kernel accepts packed cu_seqlens boundaries.
+    auto right_attention = batch.attention_mask.clone();
+    auto right_target = batch.target_mask.clone();
+    right_attention.select(1, right_attention.size(1) - 1).fill_(false);
+    right_target.select(1, right_target.size(1) - 1).fill_(0.0);
+    const double right_loss = qwen36_eval_step(
+        fixed_contexts.distributed, &batch.input_ids, &right_target,
+        &right_attention);
+    assert(std::isfinite(right_loss));
+    // A strict-right-padded token must not enter the recurrent state or affect
+    // any scored position. Changing only that token should be observationally
+    // invisible to the GDN eval path.
+    auto padded_ids_a = batch.input_ids.clone();
+    auto padded_ids_b = batch.input_ids.clone();
+    padded_ids_a.select(1, padded_ids_a.size(1) - 1).fill_(17);
+    padded_ids_b.select(1, padded_ids_b.size(1) - 1).fill_(63);
+    const double padded_loss_a = qwen36_eval_step(
+        fixed_contexts.distributed, &padded_ids_a, &right_target,
+        &right_attention);
+    const double padded_loss_b = qwen36_eval_step(
+        fixed_contexts.distributed, &padded_ids_b, &right_target,
+        &right_attention);
+    assert(std::isfinite(padded_loss_a) && std::isfinite(padded_loss_b));
+    assert(std::fabs(padded_loss_a - padded_loss_b) < 1e-6);
+    auto invalid_attention = batch.attention_mask.clone();
+    invalid_attention.index_put_({0, 0}, false);
+    const double invalid_mask_loss = qwen36_eval_step(
+        fixed_contexts.distributed, &batch.input_ids, &batch.target_mask,
+        &invalid_attention);
+    assert(invalid_mask_loss < 0.0);
+    check_fixed_path(
+        fixed_contexts, fixed_fixtures, short_batch, batch, rank);
+    qwen36_free_training_context(fixed_contexts.reference);
+    qwen36_free_training_context(fixed_contexts.distributed);
+
+    auto dynamic_contexts = create_context_pair(local_weights, full_weights,
+        local_embed, embed, final_norm, local_lm_head, lm_head,
+        configs, "in_proj_qkv", false);
+    auto dynamic_fixtures = make_lora_fixtures(rank, 4000);
+    qwen36_set_checkpoint(dynamic_contexts.distributed, 1, 1);
+    qwen36_set_checkpoint(dynamic_contexts.reference, 1, 1);
+    const int64_t build_before =
+        qwen36_get_lora_batch_projection_build_count(dynamic_contexts.distributed);
+    const int64_t upload_before =
+        qwen36_get_lora_batch_scaling_upload_count(dynamic_contexts.distributed);
+    check_dynamic_path(
+        dynamic_contexts, dynamic_fixtures, dynamic_batch, rank);
+    // The native list ABI supports a zero-capacity count query. Exercise a
+    // registry larger than the historical 64-entry Rust buffer so truncation
+    // cannot silently reappear in train_multi_lora().
+    const int64_t target_layers[kLayers] = {0, 1};
+    constexpr const char* targets =
+        "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj";
+    for (int i = 0; i < 64; ++i) {
+        assert(qwen36_add_lora(dynamic_contexts.distributed,
+            kLoraRank, kLoraRank, target_layers, kLayers, targets) > 0);
+        assert(qwen36_add_lora(dynamic_contexts.reference,
+            kLoraRank, kLoraRank, target_layers, kLayers, targets) > 0);
+    }
+    const int64_t dynamic_count = qwen36_list_lora(
+        dynamic_contexts.distributed, nullptr, 0);
+    assert(dynamic_count == 65);
+    std::vector<int64_t> listed_ids(static_cast<size_t>(dynamic_count));
+    assert(qwen36_list_lora(dynamic_contexts.distributed,
+        listed_ids.data(), dynamic_count) == dynamic_count);
+    const int64_t build_after =
+        qwen36_get_lora_batch_projection_build_count(dynamic_contexts.distributed);
+    const int64_t upload_after =
+        qwen36_get_lora_batch_scaling_upload_count(dynamic_contexts.distributed);
+    // Two GDN layers x five LoRA projections are built once in the forward
+    // and once per one-layer recompute group; scaling is uploaded once and
+    // reused by both backward groups.
+    assert(build_after - build_before == 20);
+    assert(upload_after - upload_before == 1);
+    qwen36_free_training_context(dynamic_contexts.reference);
+    qwen36_free_training_context(dynamic_contexts.distributed);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp
new file mode 100644
index 00000000..e6d6970f
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp
@@ -0,0 +1,300 @@
+#include <ATen/ATen.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_create_training_context_ex(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*, int32_t);
+extern "C" int32_t qwen36_init_nccl(void*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t);
+extern "C" int32_t qwen36_abort_gradient_accumulation(void*);
+extern "C" double qwen36_eval_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_micro_step(
+    void*, void*, void*, void*, double, int32_t);
+extern "C" int64_t qwen36_export_optimizer_state(
+    void*, void**, void**, int64_t);
+extern "C" void qwen36_free_training_context(void*);
+
+static at::Tensor deterministic(
+    std::initializer_list<int64_t> shape, double scale, int64_t offset = 0
+) {
+    int64_t count = 1;
+    for (int64_t dim : shape) count *= dim;
+    return ((at::arange(count,
+                 at::TensorOptions().device(at::kCUDA).dtype(at::kFloat))
+                 .add(offset).remainder(23) - 11.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static void append_gdn_layer(
+    std::vector<at::Tensor>& weights, int64_t hidden, int64_t intermediate,
+    int64_t state_dim, int64_t layer
+) {
+    const int64_t qkv_dim = 3 * state_dim;
+    const auto ones = at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16);
+    weights.push_back(at::ones({hidden}, ones));
+    weights.push_back(at::ones({hidden}, ones));
+    weights.push_back(deterministic({qkv_dim, hidden}, 0.0020, layer));
+    weights.push_back(deterministic({state_dim, hidden}, 0.0022, layer + 1));
+    weights.push_back(deterministic({1, hidden}, 0.0015, layer + 2));
+    weights.push_back(deterministic({1, hidden}, 0.0017, layer + 3));
+    weights.push_back(deterministic({1}, 0.0010, layer + 4));
+    weights.push_back(deterministic({1}, 0.0010, layer + 5));
+    weights.push_back(deterministic({qkv_dim, 1, 4}, 0.0012, layer + 6));
+    weights.push_back(at::ones({state_dim}, ones));
+    weights.push_back(deterministic({hidden, state_dim}, 0.0020, layer + 7));
+    weights.push_back(deterministic({intermediate, hidden}, 0.0020, layer + 8));
+    weights.push_back(deterministic({intermediate, hidden}, 0.0018, layer + 9));
+    weights.push_back(deterministic({hidden, intermediate}, 0.0020, layer + 10));
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) {
+    return (lhs - rhs).abs().max().item<double>();
+}
+
+int main() {
+    const int rank = std::atoi(std::getenv("RANK"));
+    const int world = std::atoi(std::getenv("WORLD_SIZE"));
+    assert(world == 2 && rank >= 0 && rank < world);
+    qwen36_set_cuda_device(rank);
+
+    constexpr int64_t hidden = 8;
+    constexpr int64_t intermediate = 12;
+    constexpr int64_t state_dim = 128;
+    constexpr int64_t qkv_dim = 3 * state_dim;
+    constexpr int64_t vocab = 16;
+    constexpr int64_t lora_rank = 4;
+    constexpr int64_t local_rank = lora_rank / 2;
+    constexpr int64_t slots_per_layer = 8;
+
+    std::vector<at::Tensor> weights;
+    append_gdn_layer(weights, hidden, intermediate, state_dim, 0);
+    append_gdn_layer(weights, hidden, intermediate, state_dim, 1);
+    for (auto& weight : weights) weight.set_requires_grad(false);
+    auto weight_ptrs = pointers(weights);
+
+    auto embed = deterministic({vocab, hidden}, 0.0030);
+    auto final_norm = at::ones(
+        {hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+    auto lm_head = deterministic({vocab, hidden}, 0.0025);
+    LayerConfig configs[2]{};
+    for (auto& config : configs) {
+        config.layer_type = 1;
+        config.num_k_heads = 1;
+        config.key_dim = state_dim;
+        config.num_v_heads = 1;
+        config.val_dim = state_dim;
+        config.conv_kernel = 4;
+        config.rms_eps = 1e-5;
+        config.intermediate_size = intermediate;
+    }
+    const int64_t target_layers[2] = {0, 1};
+
+    setenv("TP_SIZE", "2", 1);
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    void* distributed = qwen36_create_training_context_ex(
+        weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head,
+        configs, 2, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        // This regression exercises latent-rank-only TP with replicated GDN
+        // base weights. ABI16 bit 0 explicitly enables base GDN head TP.
+        target_layers, 2, "in_proj_qkv", 0);
+    assert(distributed && qwen36_init_nccl(distributed) == 0);
+
+    setenv("TP_SIZE", "1", 1);
+    void* reference = qwen36_create_training_context(
+        weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head,
+        configs, 2, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        target_layers, 2, "in_proj_qkv");
+    assert(reference);
+
+    std::vector<at::Tensor> full_a;
+    std::vector<at::Tensor> full_b;
+    for (int64_t layer = 0; layer < 2; ++layer) {
+        full_a.push_back(deterministic(
+            {lora_rank, hidden}, 0.0010, 20 + layer));
+        full_b.push_back(deterministic(
+            {qkv_dim, lora_rank}, 0.0008, 30 + layer));
+        auto local_a = full_a.back().narrow(
+            0, rank * local_rank, local_rank).contiguous();
+        auto local_b = full_b.back().narrow(
+            1, rank * local_rank, local_rank).contiguous();
+        const int64_t slot = layer * slots_per_layer;
+        assert(qwen36_set_lora_tensor(distributed, slot, 0, &local_a) == 0);
+        assert(qwen36_set_lora_tensor(distributed, slot, 1, &local_b) == 0);
+        assert(qwen36_set_lora_tensor(reference, slot, 0, &full_a.back()) == 0);
+        assert(qwen36_set_lora_tensor(reference, slot, 1, &full_b.back()) == 0);
+    }
+
+    auto input_ids = at::tensor({1, 2, 3},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 3});
+    auto target_mask = at::ones({1, 3},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    auto attention_mask = at::ones({1, 3},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+
+    const double distributed_eval = qwen36_eval_step(
+        distributed, &input_ids, &target_mask, &attention_mask);
+    const double reference_eval = qwen36_eval_step(
+        reference, &input_ids, &target_mask, &attention_mask);
+    assert(distributed_eval > 0.0 && reference_eval > 0.0);
+
+    const double distributed_micro = qwen36_train_micro_step(
+        distributed, &input_ids, &target_mask, &attention_mask, 1.0, 0);
+    const double reference_micro = qwen36_train_micro_step(
+        reference, &input_ids, &target_mask, &attention_mask, 1.0, 0);
+    assert(distributed_micro > 0.0 && reference_micro > 0.0);
+
+    double max_a_grad_diff = 0.0;
+    double max_b_grad_diff = 0.0;
+    for (int64_t layer = 0; layer < 2; ++layer) {
+        const int64_t slot = layer * slots_per_layer;
+        auto* local_a_grad = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_grad_accumulator(distributed, slot, 0));
+        auto* local_b_grad = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_grad_accumulator(distributed, slot, 1));
+        auto* full_a_grad = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_grad_accumulator(reference, slot, 0));
+        auto* full_b_grad = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_grad_accumulator(reference, slot, 1));
+        assert(local_a_grad && local_b_grad && full_a_grad && full_b_grad);
+        max_a_grad_diff = std::max(max_a_grad_diff, max_diff(
+            *local_a_grad, full_a_grad->narrow(
+                0, rank * local_rank, local_rank)));
+        max_b_grad_diff = std::max(max_b_grad_diff, max_diff(
+            *local_b_grad, full_b_grad->narrow(
+                1, rank * local_rank, local_rank)));
+    }
+    assert(qwen36_abort_gradient_accumulation(distributed) == 0);
+    assert(qwen36_abort_gradient_accumulation(reference) == 0);
+
+    const double distributed_loss = qwen36_train_step(
+        distributed, &input_ids, &target_mask, &attention_mask);
+    const double reference_loss = qwen36_train_step(
+        reference, &input_ids, &target_mask, &attention_mask);
+    assert(distributed_loss > 0.0 && reference_loss > 0.0);
+
+    constexpr int64_t optimizer_count = 2 * 2 * slots_per_layer;
+    std::vector<void*> local_m(optimizer_count), local_v(optimizer_count);
+    std::vector<void*> full_m(optimizer_count), full_v(optimizer_count);
+    assert(qwen36_export_optimizer_state(
+        distributed, local_m.data(), local_v.data(), optimizer_count) ==
+        optimizer_count);
+    assert(qwen36_export_optimizer_state(
+        reference, full_m.data(), full_v.data(), optimizer_count) ==
+        optimizer_count);
+
+    double max_a_param_diff = 0.0;
+    double max_b_param_diff = 0.0;
+    double max_m_diff = 0.0;
+    double max_v_diff = 0.0;
+    double max_adam_error = 0.0;
+    for (int64_t layer = 0; layer < 2; ++layer) {
+        const int64_t slot = layer * slots_per_layer;
+        auto* local_a = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(distributed, slot));
+        auto* local_b = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(distributed, slot));
+        auto* full_a_param = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_a(reference, slot));
+        auto* full_b_param = reinterpret_cast<at::Tensor*>(
+            qwen36_get_lora_b(reference, slot));
+        assert(local_a && local_b && full_a_param && full_b_param);
+        max_a_param_diff = std::max(max_a_param_diff, max_diff(
+            *local_a, full_a_param->narrow(
+                0, rank * local_rank, local_rank)));
+        max_b_param_diff = std::max(max_b_param_diff, max_diff(
+            *local_b, full_b_param->narrow(
+                1, rank * local_rank, local_rank)));
+
+        auto* local_m_a = reinterpret_cast<at::Tensor*>(local_m[2 * slot]);
+        auto* local_m_b = reinterpret_cast<at::Tensor*>(local_m[2 * slot + 1]);
+        auto* local_v_a = reinterpret_cast<at::Tensor*>(local_v[2 * slot]);
+        auto* local_v_b = reinterpret_cast<at::Tensor*>(local_v[2 * slot + 1]);
+        auto* full_m_a = reinterpret_cast<at::Tensor*>(full_m[2 * slot]);
+        auto* full_m_b = reinterpret_cast<at::Tensor*>(full_m[2 * slot + 1]);
+        auto* full_v_a = reinterpret_cast<at::Tensor*>(full_v[2 * slot]);
+        auto* full_v_b = reinterpret_cast<at::Tensor*>(full_v[2 * slot + 1]);
+        assert(local_m_a && local_m_b && local_v_a && local_v_b);
+        assert(full_m_a && full_m_b && full_v_a && full_v_b);
+        max_m_diff = std::max({max_m_diff,
+            max_diff(*local_m_a, full_m_a->narrow(
+                0, rank * local_rank, local_rank)),
+            max_diff(*local_m_b, full_m_b->narrow(
+                1, rank * local_rank, local_rank))});
+        max_v_diff = std::max({max_v_diff,
+            max_diff(*local_v_a, full_v_a->narrow(
+                0, rank * local_rank, local_rank)),
+            max_diff(*local_v_b, full_v_b->narrow(
+                1, rank * local_rank, local_rank))});
+
+        auto local_a_before = full_a[layer].narrow(
+            0, rank * local_rank, local_rank).contiguous();
+        auto local_b_before = full_b[layer].narrow(
+            1, rank * local_rank, local_rank).contiguous();
+        auto expected_a = (
+            local_a_before.to(at::kFloat) - 1e-3 *
+                (*local_m_a / (1.0 - 0.9)) /
+                (((*local_v_a / (1.0 - 0.999)).sqrt()) + 1e-8))
+            .to(at::kBFloat16);
+        auto expected_b = (
+            local_b_before.to(at::kFloat) - 1e-3 *
+                (*local_m_b / (1.0 - 0.9)) /
+                (((*local_v_b / (1.0 - 0.999)).sqrt()) + 1e-8))
+            .to(at::kBFloat16);
+        max_adam_error = std::max({max_adam_error,
+            max_diff(*local_a, expected_a), max_diff(*local_b, expected_b)});
+    }
+
+    std::printf(
+        "latent_tp_two_layer_smoke rank=%d eval_diff=%0.8e loss_diff=%0.8e a_grad_diff=%0.8e b_grad_diff=%0.8e m_diff=%0.8e v_diff=%0.8e adam_error=%0.8e a_param_diff=%0.8e b_param_diff=%0.8e\n",
+        rank, std::abs(distributed_eval - reference_eval),
+        std::abs(distributed_loss - reference_loss), max_a_grad_diff,
+        max_b_grad_diff, max_m_diff, max_v_diff, max_adam_error,
+        max_a_param_diff, max_b_param_diff);
+    std::fflush(stdout);
+    assert(std::abs(distributed_eval - reference_eval) < 5e-3);
+    assert(std::abs(distributed_loss - reference_loss) < 5e-3);
+    assert(max_a_grad_diff < 5e-4 && max_b_grad_diff < 5e-4);
+    assert(max_m_diff < 5e-5 && max_v_diff < 5e-8);
+    assert(max_adam_error == 0.0);
+    assert(max_a_param_diff < 1e-3 && max_b_param_diff < 1e-3);
+
+    qwen36_free_training_context(reference);
+    qwen36_free_training_context(distributed);
+    return 0;
+}
diff --git a/crates/rustrain-qwen3-6/tests/native_tp_mlp_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_mlp_smoke.cpp
new file mode 100644
index 00000000..52da9995
--- /dev/null
+++ b/crates/rustrain-qwen3-6/tests/native_tp_mlp_smoke.cpp
@@ -0,0 +1,225 @@
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+struct LayerConfig {
+    int64_t layer_type, num_heads, num_kv_heads, head_dim;
+    int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel;
+    double partial_rotary_factor, rope_theta, rms_eps;
+    int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count;
+    int64_t intermediate_size;
+    int32_t norm_topk_prob;
+    void* nccl_comm;
+    void* nccl_stream;
+};
+
+extern "C" void qwen36_set_cuda_device(int32_t);
+extern "C" void* qwen36_create_training_context(
+    void**, int64_t, void*, void*, void*, void*, int64_t, int32_t,
+    double, double, double, double, double, int64_t, double, int64_t,
+    const int64_t*, int64_t, const char*);
+extern "C" int32_t qwen36_init_nccl(void*);
+extern "C" int32_t qwen36_set_base_tp_mlp(void*, int32_t);
+extern "C" int32_t qwen36_set_mtp_weights(
+    void*, void*, void*, void*, void*, void**, int64_t, void*, int64_t);
+extern "C" int64_t qwen36_add_lora(
+    void*, int64_t, double, const int64_t*, int64_t, const char*);
+extern "C" void* qwen36_get_lora_a(void*, int64_t);
+extern "C" void* qwen36_get_lora_b(void*, int64_t);
+extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*);
+extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t);
+extern "C" int32_t qwen36_abort_gradient_accumulation(void*);
+extern "C" double qwen36_eval_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_step(void*, void*, void*, void*);
+extern "C" double qwen36_train_micro_step(
+    void*, void*, void*, void*, double, int32_t);
+extern "C" void qwen36_free_training_context(void*);
+
+static at::Tensor deterministic(std::initializer_list<int64_t> shape, double scale) {
+    int64_t count = 1;
+    for (int64_t dim : shape) count *= dim;
+    return ((at::arange(count,
+                 at::TensorOptions().device(at::kCUDA).dtype(at::kFloat))
+                 .remainder(17) - 8.0) * scale)
+        .reshape(shape).to(at::kBFloat16);
+}
+
+static std::vector<void*> pointers(std::vector<at::Tensor>& tensors) {
+    std::vector<void*> result;
+    result.reserve(tensors.size());
+    for (auto& tensor : tensors) result.push_back(&tensor);
+    return result;
+}
+
+int main() {
+    const int process_rank = std::atoi(std::getenv("RANK"));
+    const int world = std::atoi(std::getenv("WORLD_SIZE"));
+    assert(world == 2 && process_rank >= 0 && process_rank < world);
+    qwen36_set_cuda_device(process_rank);
+
+    constexpr int64_t hidden = 8;
+    constexpr int64_t intermediate = 12;
+    constexpr int64_t vocab = 16;
+    constexpr int64_t lora_rank = 4;
+    constexpr int64_t local_lora_rank = lora_rank / 2;
+
+    std::vector<at::Tensor> full_weights;
+    full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({2 * hidden, hidden}, 0.01));
+    full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({hidden, hidden}, 0.012));
+    full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)));
+    full_weights.push_back(deterministic({hidden, hidden}, 0.008));
+    full_weights.push_back(deterministic({hidden, hidden}, 0.011));
+    full_weights.push_back(deterministic({intermediate, hidden}, 0.009));
+    full_weights.push_back(deterministic({intermediate, hidden}, 0.007));
+    full_weights.push_back(deterministic({hidden, intermediate}, 0.01));
+    for (auto& weight : full_weights) weight.set_requires_grad(false);
+
+    const int64_t intermediate_start = process_rank * (intermediate / world);
+    std::vector<at::Tensor> local_weights(full_weights.begin(), full_weights.begin() + 8);
+    local_weights.push_back(full_weights[8].narrow(0, intermediate_start, intermediate / world).contiguous());
+    local_weights.push_back(full_weights[9].narrow(0, intermediate_start, intermediate / world).contiguous());
+    local_weights.push_back(full_weights[10].narrow(1, intermediate_start, intermediate / world).contiguous());
+    assert(local_weights[8].sizes() == at::IntArrayRef({intermediate / 2, hidden}));
+    assert(local_weights[9].sizes() == at::IntArrayRef({intermediate / 2, hidden}));
+    assert(local_weights[10].sizes() == at::IntArrayRef({hidden, intermediate / 2}));
+
+    auto embed = deterministic({vocab, hidden}, 0.02);
+    auto final_norm = at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16));
+    auto lm_head = deterministic({vocab, hidden}, 0.015);
+    LayerConfig config{};
+    config.layer_type = 0;
+    config.num_heads = 1;
+    config.num_kv_heads = 1;
+    config.head_dim = hidden;
+    config.partial_rotary_factor = 1.0;
+    config.rope_theta = 10000.0;
+    config.rms_eps = 1e-5;
+    config.intermediate_size = intermediate;
+
+    const int64_t target_layer = 0;
+    auto local_ptrs = pointers(local_weights);
+    setenv("TP_SIZE", "2", 1);
+    unsetenv("RUSTRAIN_DATA_PARALLEL");
+    void* rejected = qwen36_create_training_context(
+        local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, "gate_proj");
+    assert(rejected && qwen36_set_base_tp_mlp(rejected, 1) != 0);
+    qwen36_free_training_context(rejected);
+
+    void* dynamic_rejected = qwen36_create_training_context(
+        local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, "q_proj");
+    assert(dynamic_rejected);
+    assert(qwen36_add_lora(
+        dynamic_rejected, lora_rank, 1.0, &target_layer, 1, "gate_proj") > 0);
+    assert(qwen36_set_base_tp_mlp(dynamic_rejected, 1) != 0);
+    qwen36_free_training_context(dynamic_rejected);
+
+    void* distributed = qwen36_create_training_context(
+        local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, "q_proj");
+    assert(distributed && qwen36_set_base_tp_mlp(distributed, 1) == 0);
+    assert(qwen36_set_base_tp_mlp(distributed, 0) != 0);
+    assert(qwen36_set_mtp_weights(
+        distributed, nullptr, nullptr, nullptr, nullptr, nullptr, 0, nullptr, 0) != 0);
+    assert(qwen36_add_lora(
+        distributed, lora_rank, 1.0, &target_layer, 1, nullptr) < 0);
+    assert(qwen36_init_nccl(distributed) == 0);
+
+    setenv("TP_SIZE", "1", 1);
+    auto full_ptrs = pointers(full_weights);
+    void* reference = qwen36_create_training_context(
+        full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head,
+        &config, 1, static_cast<int32_t>(at::kBFloat16),
+        1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank,
+        &target_layer, 1, "q_proj");
+    assert(reference);
+
+    auto full_a = deterministic({lora_rank, hidden}, 0.002);
+    auto full_b = deterministic({2 * hidden, lora_rank}, 0.001);
+    auto local_a = full_a.narrow(0, process_rank * local_lora_rank, local_lora_rank).contiguous();
+    auto local_b = full_b.narrow(1, process_rank * local_lora_rank, local_lora_rank).contiguous();
+    assert(qwen36_set_lora_tensor(distributed, 0, 0, &local_a) == 0);
+    assert(qwen36_set_lora_tensor(distributed, 0, 1, &local_b) == 0);
+    assert(qwen36_set_lora_tensor(reference, 0, 0, &full_a) == 0);
+    assert(qwen36_set_lora_tensor(reference, 0, 1, &full_b) == 0);
+    auto local_a_before = local_a.clone();
+    auto local_b_before = local_b.clone();
+
+    auto input_ids = at::tensor({1, 2, 3},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 3});
+    auto target_mask = at::ones({1, 3},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));
+    auto attention_mask = at::ones({1, 3},
+        at::TensorOptions().device(at::kCUDA).dtype(at::kBool));
+
+    const double distributed_eval = qwen36_eval_step(
+        distributed, &input_ids, &target_mask, &attention_mask);
+    const double reference_eval = qwen36_eval_step(
+        reference, &input_ids, &target_mask, &attention_mask);
+    assert(distributed_eval > 0.0 && reference_eval > 0.0);
+    assert(std::abs(distributed_eval - reference_eval) < 1e-4);
+
+    const double distributed_micro = qwen36_train_micro_step(
+        distributed, &input_ids, &target_mask, &attention_mask, 1.0, 0);
+    const double reference_micro = qwen36_train_micro_step(
+        reference, &input_ids, &target_mask, &attention_mask, 1.0, 0);
+    assert(distributed_micro > 0.0 && reference_micro > 0.0);
+    auto* distributed_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(distributed, 0, 1));
+    auto* reference_b_grad = reinterpret_cast<at::Tensor*>(
+        qwen36_get_lora_grad_accumulator(reference, 0, 1));
+    assert(distributed_b_grad && reference_b_grad);
+    auto reference_b_grad_slice = reference_b_grad->narrow(
+        1, process_rank * local_lora_rank, local_lora_rank);
+    const double b_grad_diff =
+        (*distributed_b_grad - reference_b_grad_slice).abs().max().item<double>();
+    assert(distributed_b_grad->abs().max().item<double>() > 0.0);
+    assert(b_grad_diff < 1e-4);
+    assert(qwen36_abort_gradient_accumulation(distributed) == 0);
+    assert(qwen36_abort_gradient_accumulation(reference) == 0);
+
+    const double distributed_loss = qwen36_train_step(
+        distributed, &input_ids, &target_mask, &attention_mask);
+    const double reference_loss = qwen36_train_step(
+        reference, &input_ids, &target_mask, &attention_mask);
+    assert(distributed_loss > 0.0 && reference_loss > 0.0);
+    assert(std::abs(distributed_loss - reference_loss) < 1e-4);
+
+    auto* updated_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(distributed, 0));
+    auto* updated_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(distributed, 0));
+    auto* reference_a = reinterpret_cast<at::Tensor*>(qwen36_get_lora_a(reference, 0));
+    auto* reference_b = reinterpret_cast<at::Tensor*>(qwen36_get_lora_b(reference, 0));
+    auto reference_a_slice = reference_a->narrow(
+        0, process_rank * local_lora_rank, local_lora_rank);
+    auto reference_b_slice = reference_b->narrow(
+        1, process_rank * local_lora_rank, local_lora_rank);
+    const double a_diff = (*updated_a - reference_a_slice).abs().max().item<double>();
+    const double b_diff = (*updated_b - reference_b_slice).abs().max().item<double>();
+    const double a_update = (*updated_a - local_a_before).abs().max().item<double>();
+    const double b_update = (*updated_b - local_b_before).abs().max().item<double>();
+    std::printf(
+        "base_tp_mlp_smoke rank=%d eval_diff=%0.8e loss_diff=%0.8e b_grad_diff=%0.8e a_diff=%0.8e b_diff=%0.8e\n",
+        process_rank, std::abs(distributed_eval - reference_eval),
+        std::abs(distributed_loss - reference_loss), b_grad_diff, a_diff, b_diff);
+    assert(a_diff < 1e-5 && b_diff < 1e-5);
+    assert(a_update > 0.0 || b_update > 0.0);
+
+    qwen36_free_training_context(reference);
+    qwen36_free_training_context(distributed);
+    return 0;
+}
diff --git a/crates/rustrain-server/Cargo.toml b/crates/rustrain-server/Cargo.toml
index cf36ea5a..fdb2da4f 100644
--- a/crates/rustrain-server/Cargo.toml
+++ b/crates/rustrain-server/Cargo.toml
@@ -9,6 +9,7 @@ rustrain-qwen3-6 = { path = "../rustrain-qwen3-6" }
 rustrain-checkpoint = { path = "../rustrain-checkpoint" }
 rustrain-train = { path = "../rustrain-train" }
 rustrain-ipc = { path = "../rustrain-ipc" }
+rustrain-parallel = { path = "../rustrain-parallel" }
 anyhow.workspace = true
 axum.workspace = true
 base64 = "0.22"
@@ -30,3 +31,6 @@ tracing.workspace = true
 
 [build-dependencies]
 tonic-build.workspace = true
+
+[dev-dependencies]
+tempfile.workspace = true
diff --git a/crates/rustrain-server/proto/train.proto b/crates/rustrain-server/proto/train.proto
index 644f72c1..839ae6df 100644
--- a/crates/rustrain-server/proto/train.proto
+++ b/crates/rustrain-server/proto/train.proto
@@ -10,6 +10,7 @@ service TrainService {
   rpc SaveCheckpoint(SaveCheckpointRequest) returns (CheckpointInfo);
   rpc LoadCheckpoint(LoadCheckpointRequest) returns (CheckpointInfo);
   rpc ExportAdapter(ExportAdapterRequest) returns (AdapterInfo);
+  rpc ImportAdapter(ImportAdapterRequest) returns (ImportAdapterResponse);
   rpc StreamMetrics(StreamMetricsRequest) returns (stream StepMetric);
   rpc GetStatus(GetStatusRequest) returns (SessionStatus);
   rpc AddLoRA(AddLoRARequest) returns (AddLoRAResponse);
@@ -94,12 +95,21 @@ message CheckpointInfo {
 message ExportAdapterRequest {
   string session_id = 1;
   string path = 2;
+  int64 adapter_id = 3; // 0 exports the fixed adapter
 }
 message AdapterInfo {
   string path = 1;
   int64 param_count = 2;
 }
 
+message ImportAdapterRequest {
+  string session_id = 1;
+  string path = 2;
+}
+message ImportAdapterResponse {
+  int64 adapter_id = 1;
+}
+
 message StreamMetricsRequest {
   string session_id = 1;
 }
@@ -121,12 +131,20 @@ message SessionStatus {
   string model_path = 4;
 }
 
+message OptionalDouble {
+  double value = 1;
+}
+
 message AddLoRARequest {
   string session_id = 1;
   int64 rank = 2;
   double alpha = 3;
   repeated int64 target_layers = 4;
   string target_modules = 5;  // comma-separated, empty = all
+  OptionalDouble optimizer_lr = 6;  // omitted inherits the session optimizer
+  OptionalDouble optimizer_beta1 = 7;
+  OptionalDouble optimizer_beta2 = 8;
+  OptionalDouble optimizer_eps = 9;
 }
 message AddLoRAResponse {
   int64 adapter_id = 1;
diff --git a/crates/rustrain-server/src/api.rs b/crates/rustrain-server/src/api.rs
index 3ed5b008..b65a5062 100644
--- a/crates/rustrain-server/src/api.rs
+++ b/crates/rustrain-server/src/api.rs
@@ -1,18 +1,27 @@
 //! HTTP API (axum) — RESTful endpoints for training session management.
 
 use axum::{
+    body::Bytes,
+    extract::DefaultBodyLimit,
+    extract::Request,
     extract::{Path, State},
-    http::StatusCode,
+    http::{HeaderMap, HeaderValue, StatusCode},
+    middleware::{self, Next},
     response::sse::{Event, KeepAlive, Sse},
+    response::{IntoResponse, Response},
     routing::{get, post},
     Json, Router,
 };
 use futures::stream::{self, Stream};
 use serde::{Deserialize, Serialize};
+use std::collections::{HashMap, HashSet, VecDeque};
 use std::convert::Infallible;
-use std::sync::Arc;
+use std::sync::{Arc, Condvar, Mutex};
+use std::time::{Duration, Instant};
+use tokio::sync::{oneshot, Semaphore};
 use tokio_stream::StreamExt;
 
+use crate::ep_dispatch::{configured_queue_capacity, EpDispatchScheduleError, EpDispatchScheduler};
 use crate::metrics::StepMetric;
 use crate::session::{InitLoRARequest, SessLoadDatasetRequest, SessLoadModelRequest, TrainInput};
 use crate::state::SessionManager;
@@ -24,6 +33,500 @@ pub struct AppState {
 /// EP mode: HTTP server dispatches to workers via IPC coordinator.
 pub struct EpAppState {
     pub coordinator: Arc<crate::ep::EpCoordinator>,
+    pub world_size: usize,
+}
+
+struct EpRouterState {
+    coordinator: Arc<crate::ep::EpCoordinator>,
+    world_size: usize,
+    dispatcher: EpDispatchScheduler,
+    dispatch_submission: Mutex<()>,
+    multi_lora_batcher: MultiLoraBatcher,
+}
+
+const DEFAULT_MULTI_LORA_BATCH_WINDOW_US: u64 = 2_000;
+const HARD_MAX_MULTI_LORA_BATCH_WINDOW_US: u64 = 100_000;
+const DEFAULT_MULTI_LORA_BATCH_REQUESTS: usize = 16;
+const HARD_MAX_MULTI_LORA_BATCH_REQUESTS: usize = 4_096;
+const DEFAULT_MULTI_LORA_BATCH_OPEN_WINDOWS: usize = 8;
+const HARD_MAX_MULTI_LORA_BATCH_OPEN_WINDOWS: usize = 64;
+const DEFAULT_MULTI_LORA_BATCH_ADAPTERS: usize = 64;
+const DEFAULT_MULTI_LORA_BATCH_RANK_WORK: usize = 32_768;
+const HARD_MAX_MULTI_LORA_BATCH_RANK_WORK: usize = 4_194_304;
+// AdapterLoss is JSON-encoded into a fixed 256 KiB IPC result slot. Keep a
+// conservative margin for worst-case i64/f64 strings and the result envelope.
+const HARD_MAX_MULTI_LORA_BATCH_ADAPTERS: usize = 2_048;
+const MULTI_LORA_CAPABILITY_HEADER: &str = "x-rustrain-multi-lora-capability";
+const MULTI_LORA_CAPABILITY_V1: &str = "v1";
+const MULTI_LORA_RESPONSE_CAPABILITIES: &[&str] = &[
+    "per_adapter_loss_v1",
+    "optimizer_steps_v1",
+    "coalesced_loss_scope_v1",
+];
+
+#[derive(Clone, Copy)]
+struct MultiLoraBatchConfig {
+    window: Duration,
+    max_requests: usize,
+    max_open_windows: usize,
+    max_adapters: usize,
+    max_rank_work: usize,
+    max_payload_bytes: usize,
+}
+
+impl MultiLoraBatchConfig {
+    fn from_env() -> Self {
+        let window_us = configured_positive_us(
+            "RUSTRAIN_MULTI_LORA_BATCH_WINDOW_US",
+            DEFAULT_MULTI_LORA_BATCH_WINDOW_US,
+            HARD_MAX_MULTI_LORA_BATCH_WINDOW_US,
+        );
+        let max_requests = configured_positive_usize(
+            "RUSTRAIN_MULTI_LORA_BATCH_MAX_REQUESTS",
+            DEFAULT_MULTI_LORA_BATCH_REQUESTS,
+            HARD_MAX_MULTI_LORA_BATCH_REQUESTS,
+        );
+        let max_open_windows = configured_positive_usize(
+            "RUSTRAIN_MULTI_LORA_BATCH_MAX_OPEN_WINDOWS",
+            DEFAULT_MULTI_LORA_BATCH_OPEN_WINDOWS,
+            HARD_MAX_MULTI_LORA_BATCH_OPEN_WINDOWS,
+        );
+        let max_adapters = configured_positive_usize(
+            "RUSTRAIN_MULTI_LORA_BATCH_MAX_ADAPTERS",
+            DEFAULT_MULTI_LORA_BATCH_ADAPTERS,
+            HARD_MAX_MULTI_LORA_BATCH_ADAPTERS,
+        );
+        let max_rank_work = configured_positive_usize(
+            "RUSTRAIN_MULTI_LORA_BATCH_MAX_RANK_WORK",
+            DEFAULT_MULTI_LORA_BATCH_RANK_WORK,
+            HARD_MAX_MULTI_LORA_BATCH_RANK_WORK,
+        );
+        let slab_bytes = configured_positive_usize(
+            "RUSTRAIN_EP_TENSOR_SLAB_BYTES",
+            rustrain_ipc::DEFAULT_TENSOR_SLAB_BYTES,
+            usize::MAX,
+        );
+        let max_payload_bytes = configured_positive_usize(
+            "RUSTRAIN_MULTI_LORA_BATCH_MAX_BYTES",
+            slab_bytes,
+            slab_bytes,
+        );
+        Self {
+            window: Duration::from_micros(window_us),
+            max_requests,
+            max_open_windows,
+            max_adapters,
+            max_rank_work,
+            max_payload_bytes,
+        }
+    }
+}
+
+fn configured_positive_us(name: &str, default: u64, hard_max: u64) -> u64 {
+    std::env::var(name)
+        .ok()
+        .and_then(|value| value.parse::<u64>().ok())
+        .filter(|value| *value > 0)
+        .unwrap_or(default)
+        .min(hard_max)
+}
+
+fn configured_positive_usize(name: &str, default: usize, hard_max: usize) -> usize {
+    std::env::var(name)
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .filter(|value| *value > 0)
+        .unwrap_or(default)
+        .min(hard_max)
+}
+
+struct MultiLoraDispatchRequest {
+    session_id: String,
+    tensors: rustrain_ipc::TensorSlabRef,
+    payload: Vec<u8>,
+    n_total: usize,
+    lora_rank: i32,
+    adapter_ids: Vec<i64>,
+    expected_steps: Vec<u64>,
+    source_count: usize,
+    normalized_payload_bytes: usize,
+    allow_aggregate_loss: bool,
+    response: oneshot::Sender<MultiLoraDispatchOutcome>,
+}
+
+impl MultiLoraDispatchRequest {
+    fn coalescible(&self) -> bool {
+        // Optimistic step checks are request-local failure boundaries. Keeping
+        // them out of shared windows prevents one stale retry from failing
+        // unrelated tenants in the same flattened native command.
+        self.allow_aggregate_loss && !self.adapter_ids.is_empty() &&
+            self.expected_steps.is_empty()
+    }
+}
+
+struct MultiLoraDispatchOutcome {
+    result: rustrain_ipc::EpResult,
+    request_count: usize,
+}
+
+struct MultiLoraResponseTarget {
+    response: oneshot::Sender<MultiLoraDispatchOutcome>,
+    adapter_range: std::ops::Range<usize>,
+}
+
+fn project_multi_lora_result(
+    result: &rustrain_ipc::EpResult,
+    adapter_range: std::ops::Range<usize>,
+) -> rustrain_ipc::EpResult {
+    match result {
+        rustrain_ipc::EpResult::MultiLoraTrain {
+            loss,
+            step,
+            adapter_losses,
+            ..
+        } if adapter_losses.is_empty() => rustrain_ipc::EpResult::Train {
+            loss: *loss,
+            step: *step,
+        },
+        rustrain_ipc::EpResult::MultiLoraTrain {
+            loss,
+            step,
+            adapter_losses,
+            adapter_steps,
+        } if adapter_range.end <= adapter_losses.len()
+            && (adapter_steps.is_empty() || adapter_range.end <= adapter_steps.len()) =>
+        {
+            rustrain_ipc::EpResult::MultiLoraTrain {
+                loss: *loss,
+                step: *step,
+                adapter_losses: adapter_losses[adapter_range.clone()].to_vec(),
+                adapter_steps: if adapter_steps.is_empty() {
+                    Vec::new()
+                } else {
+                    adapter_steps[adapter_range].to_vec()
+                },
+            }
+        }
+        rustrain_ipc::EpResult::MultiLoraTrain {
+            adapter_losses,
+            adapter_steps,
+            ..
+        } => rustrain_ipc::EpResult::Error(format!(
+            "native adapter result counts (losses={}, steps={}) do not cover coalesced range {:?}",
+            adapter_losses.len(),
+            adapter_steps.len(),
+            adapter_range
+        )),
+        _ => result.clone(),
+    }
+}
+
+struct MultiLoraWindow {
+    session_id: String,
+    seq_len: usize,
+    source_count: usize,
+    adapter_ids: HashSet<i64>,
+    adapter_count: usize,
+    max_lora_rank: usize,
+    rank_work: usize,
+    payload_bytes: usize,
+    coalescible: bool,
+    validates_steps: bool,
+    requests: Vec<MultiLoraDispatchRequest>,
+}
+
+impl MultiLoraWindow {
+    fn new(request: MultiLoraDispatchRequest) -> Self {
+        let adapter_ids = request.adapter_ids.iter().copied().collect();
+        Self {
+            session_id: request.session_id.clone(),
+            seq_len: request.tensors.seq_len,
+            source_count: request.source_count,
+            adapter_ids,
+            adapter_count: request.n_total,
+            max_lora_rank: request.lora_rank.max(0) as usize,
+            rank_work: request
+                .n_total
+                .saturating_mul(request.lora_rank.max(0) as usize),
+            payload_bytes: request.normalized_payload_bytes,
+            coalescible: request.coalescible(),
+            validates_steps: !request.expected_steps.is_empty(),
+            requests: vec![request],
+        }
+    }
+
+    fn can_accept(&self, request: &MultiLoraDispatchRequest, config: MultiLoraBatchConfig) -> bool {
+        if !self.coalescible
+            || !request.coalescible()
+            || self.session_id != request.session_id
+            || self.seq_len != request.tensors.seq_len
+            || self.source_count != request.source_count
+            || self.validates_steps != !request.expected_steps.is_empty()
+            || self.requests.len() >= config.max_requests
+            || self.adapter_count.saturating_add(request.n_total) > config.max_adapters
+            || request.lora_rank <= 0
+            || self
+                .payload_bytes
+                .saturating_add(request.normalized_payload_bytes)
+                > config.max_payload_bytes
+        {
+            return false;
+        }
+        let projected_count = self.adapter_count.saturating_add(request.n_total);
+        let projected_max_rank = self.max_lora_rank.max(request.lora_rank as usize);
+        let Some(projected_rank_work) = projected_count.checked_mul(projected_max_rank) else {
+            return false;
+        };
+        if projected_rank_work > config.max_rank_work {
+            return false;
+        }
+        request
+            .adapter_ids
+            .iter()
+            .all(|adapter_id| !self.adapter_ids.contains(adapter_id))
+    }
+
+    fn push(&mut self, request: MultiLoraDispatchRequest) {
+        self.adapter_count += request.n_total;
+        self.max_lora_rank = self.max_lora_rank.max(request.lora_rank as usize);
+        self.rank_work = self.adapter_count.saturating_mul(self.max_lora_rank);
+        self.payload_bytes += request.normalized_payload_bytes;
+        self.adapter_ids.extend(request.adapter_ids.iter().copied());
+        self.requests.push(request);
+    }
+
+    fn at_capacity(&self, config: MultiLoraBatchConfig) -> bool {
+        self.requests.len() >= config.max_requests
+            || self.adapter_count >= config.max_adapters
+            || self.rank_work >= config.max_rank_work
+            || self.payload_bytes >= config.max_payload_bytes
+    }
+}
+
+#[derive(Default)]
+struct MultiLoraBatchState {
+    next_window_id: u64,
+    open_window_ids: VecDeque<u64>,
+    windows: HashMap<u64, MultiLoraWindow>,
+}
+
+#[derive(Clone)]
+struct MultiLoraBatcher {
+    state: Arc<Mutex<MultiLoraBatchState>>,
+    window_wakeup: Arc<Condvar>,
+    config: MultiLoraBatchConfig,
+}
+
+impl MultiLoraBatcher {
+    fn new(config: MultiLoraBatchConfig) -> Self {
+        Self {
+            state: Arc::new(Mutex::new(MultiLoraBatchState::default())),
+            window_wakeup: Arc::new(Condvar::new()),
+            config,
+        }
+    }
+
+    fn seal_current(&self) {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        state.open_window_ids.clear();
+        self.window_wakeup.notify_all();
+    }
+
+    fn seal_window(state: &mut MultiLoraBatchState, window_id: u64) -> bool {
+        let Some(index) = state
+            .open_window_ids
+            .iter()
+            .position(|candidate| *candidate == window_id)
+        else {
+            return false;
+        };
+        state.open_window_ids.remove(index);
+        true
+    }
+
+    fn compatible_window_id(
+        state: &MultiLoraBatchState,
+        request: &MultiLoraDispatchRequest,
+        config: MultiLoraBatchConfig,
+    ) -> Option<u64> {
+        state
+            .open_window_ids
+            .iter()
+            .rev()
+            .copied()
+            .find(|window_id| {
+                state
+                    .windows
+                    .get(window_id)
+                    .is_some_and(|window| window.can_accept(request, config))
+            })
+    }
+
+    /// Wait for a coalescing window to be sealed or for its normal deadline.
+    /// The window ID predicate makes spurious notifications harmless and
+    /// treats removal from the bounded open-window set as sealing. The
+    /// scheduler remains single-consumer, so this only changes when the
+    /// existing FIFO job becomes dispatchable.
+    fn wait_for_window(&self, window_id: u64, deadline: Instant) {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        while state.open_window_ids.contains(&window_id) {
+            let remaining = deadline.saturating_duration_since(Instant::now());
+            if remaining.is_zero() {
+                break;
+            }
+            let (next_state, timeout) = self
+                .window_wakeup
+                .wait_timeout(state, remaining)
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            state = next_state;
+            if timeout.timed_out() {
+                break;
+            }
+        }
+        Self::seal_window(&mut state, window_id);
+    }
+
+    fn seal_window_if_at_capacity(&self, state: &mut MultiLoraBatchState, window_id: u64) {
+        if state
+            .windows
+            .get(&window_id)
+            .is_some_and(|window| window.at_capacity(self.config))
+            && Self::seal_window(state, window_id)
+        {
+            self.window_wakeup.notify_all();
+        }
+    }
+
+    fn submit(
+        &self,
+        scheduler: &EpDispatchScheduler,
+        coordinator: Arc<crate::ep::EpCoordinator>,
+        request: MultiLoraDispatchRequest,
+    ) -> Result<(), EpDispatchScheduleError> {
+        if request.n_total > self.config.max_adapters
+            || request.normalized_payload_bytes > self.config.max_payload_bytes
+            || request.lora_rank <= 0
+            || request
+                .n_total
+                .checked_mul(request.lora_rank as usize)
+                .map_or(true, |rank_work| rank_work > self.config.max_rank_work)
+        {
+            return Err(EpDispatchScheduleError::QueueFull);
+        }
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        if let Some(window_id) = Self::compatible_window_id(&state, &request, self.config) {
+            state.windows.get_mut(&window_id).unwrap().push(request);
+            self.seal_window_if_at_capacity(&mut state, window_id);
+            return Ok(());
+        }
+
+        let window_id = state.next_window_id;
+        state.next_window_id = state.next_window_id.wrapping_add(1);
+        let coalescible = request.coalescible();
+        state
+            .windows
+            .insert(window_id, MultiLoraWindow::new(request));
+        if coalescible {
+            if state.open_window_ids.len() >= self.config.max_open_windows {
+                state.open_window_ids.pop_front();
+                self.window_wakeup.notify_all();
+            }
+            state.open_window_ids.push_back(window_id);
+            self.seal_window_if_at_capacity(&mut state, window_id);
+        }
+
+        let batcher = self.clone();
+        // Start the batching window at admission. The FIFO dispatcher may still be
+        // running the previous GPU step, and that compute time should consume this
+        // deadline instead of creating another idle batching delay afterwards.
+        let deadline = Instant::now()
+            + if coalescible {
+                self.config.window
+            } else {
+                Duration::ZERO
+            };
+        let scheduled = scheduler.submit(move || {
+            batcher.wait_for_window(window_id, deadline);
+            batcher.execute_window(window_id, &coordinator);
+        });
+        if let Err(error) = scheduled {
+            state.windows.remove(&window_id);
+            if Self::seal_window(&mut state, window_id) {
+                self.window_wakeup.notify_all();
+            }
+            return Err(error);
+        }
+        drop(scheduled);
+        Ok(())
+    }
+
+    fn execute_window(&self, window_id: u64, coordinator: &crate::ep::EpCoordinator) {
+        let window = {
+            let mut state = self
+                .state
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            Self::seal_window(&mut state, window_id);
+            state.windows.remove(&window_id)
+        };
+        let Some(window) = window else {
+            return;
+        };
+        let (command, payload, responses) = match build_multi_lora_window(window) {
+            Ok(result) => result,
+            Err((error, responses)) => {
+                let request_count = responses.len();
+                for target in responses {
+                    let _ = target.response.send(MultiLoraDispatchOutcome {
+                        result: rustrain_ipc::EpResult::Error(error.clone()),
+                        request_count,
+                    });
+                }
+                return;
+            }
+        };
+        if payload.len() > self.config.max_payload_bytes {
+            let request_count = responses.len();
+            for target in responses {
+                let _ = target.response.send(MultiLoraDispatchOutcome {
+                    result: rustrain_ipc::EpResult::Error(format!(
+                        "coalesced tensor payload {} exceeds configured limit {}",
+                        payload.len(),
+                        self.config.max_payload_bytes
+                    )),
+                    request_count,
+                });
+            }
+            return;
+        }
+        tracing::debug!(
+            requests = responses.len(),
+            payload_bytes = payload.len(),
+            "dispatching coalesced multi-LoRA training batch"
+        );
+        let result = if coordinator.is_healthy() {
+            coordinator.dispatch_with_slab(&command, &payload)
+        } else {
+            rustrain_ipc::EpResult::Error("EP coordinator is unavailable".to_string())
+        };
+        let request_count = responses.len();
+        for target in responses {
+            let projected = project_multi_lora_result(&result, target.adapter_range);
+            let _ = target.response.send(MultiLoraDispatchOutcome {
+                result: projected,
+                request_count,
+            });
+        }
+    }
 }
 
 pub fn router(state: Arc<AppState>) -> Router {
@@ -36,9 +539,11 @@ pub fn router(state: Arc<AppState>) -> Router {
         .route("/v1/sessions/{id}/init_lora", post(init_lora))
         .route("/v1/sessions/{id}/train_step", post(train_step))
         .route("/v1/sessions/{id}/eval_step", post(eval_step))
+        .route("/v1/sessions/{id}/eval_multi_lora", post(eval_multi_lora))
         .route("/v1/sessions/{id}/save_checkpoint", post(save_checkpoint))
         .route("/v1/sessions/{id}/load_checkpoint", post(load_checkpoint))
         .route("/v1/sessions/{id}/export_adapter", post(export_adapter))
+        .route("/v1/sessions/{id}/import_adapter", post(import_adapter))
         .route("/v1/sessions/{id}/add_lora", post(add_lora))
         .route("/v1/sessions/{id}/remove_lora", post(remove_lora))
         .route("/v1/sessions/{id}/list_lora", get(list_lora))
@@ -55,7 +560,9 @@ struct ErrorResponse {
 fn err_resp(msg: &str) -> (StatusCode, Json<ErrorResponse>) {
     (
         StatusCode::BAD_REQUEST,
-        Json(ErrorResponse { error: msg.to_string() }),
+        Json(ErrorResponse {
+            error: msg.to_string(),
+        }),
     )
 }
 
@@ -76,7 +583,11 @@ async fn delete_session(
     State(state): State<Arc<AppState>>,
     axum::extract::Path(id): axum::extract::Path<String>,
 ) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
-    state.manager.delete_session(&id).await.map_err(|e| err_resp(&e))?;
+    state
+        .manager
+        .delete_session(&id)
+        .await
+        .map_err(|e| err_resp(&e))?;
     Ok(Json(serde_json::json!({"deleted": id})))
 }
 
@@ -89,9 +600,7 @@ struct CreateSessionResponse {
     session_id: String,
 }
 
-async fn list_sessions(
-    State(state): State<Arc<AppState>>,
-) -> Json<Vec<String>> {
+async fn list_sessions(State(state): State<Arc<AppState>>) -> Json<Vec<String>> {
     Json(state.manager.list_sessions().await)
 }
 
@@ -181,7 +690,7 @@ async fn init_lora(
 #[derive(Deserialize)]
 struct InitLoRAHttp {
     rank: i64,
-    alpha: i64,
+    alpha: f64,
     target_layers: Vec<usize>,
     target_modules: Vec<String>,
     lr: f64,
@@ -204,12 +713,9 @@ async fn train_step(
         .get_session(&id)
         .await
         .ok_or_else(|| err_resp("session not found"))?;
-    let input_ids = decode_tensor(&req.input_ids)
-        .map_err(|e| err_resp(&e))?;
-    let target_mask = decode_tensor(&req.target_mask)
-        .map_err(|e| err_resp(&e))?;
-    let attention_mask = decode_tensor(&req.attention_mask)
-        .map_err(|e| err_resp(&e))?;
+    let input_ids = decode_tensor(&req.input_ids).map_err(|e| err_resp(&e))?;
+    let target_mask = decode_tensor(&req.target_mask).map_err(|e| err_resp(&e))?;
+    let attention_mask = decode_tensor(&req.attention_mask).map_err(|e| err_resp(&e))?;
 
     let mut s = session.lock().await;
     let result = s
@@ -238,6 +744,14 @@ struct TrainStepHttp {
     attention_mask: TensorHttp,
 }
 
+#[derive(Deserialize)]
+struct EvalMultiLoraHttp {
+    input_ids: TensorHttp,
+    target_mask: TensorHttp,
+    attention_mask: TensorHttp,
+    adapter_ids: Vec<i64>,
+}
+
 #[derive(Deserialize)]
 struct TrainMultiLoraHttp {
     input_ids: TensorHttp,
@@ -245,6 +759,219 @@ struct TrainMultiLoraHttp {
     attention_mask: TensorHttp,
     n_total: i32,
     lora_rank: i32,
+    #[serde(default)]
+    adapter_ids: Vec<i64>,
+    #[serde(default)]
+    expected_steps: Vec<u64>,
+    #[serde(default)]
+    allow_aggregate_loss: bool,
+}
+
+const MULTI_LORA_BINARY_MAGIC: [u8; 4] = *b"RLM1";
+const MULTI_LORA_BINARY_VERSION: u16 = 1;
+const MULTI_LORA_BINARY_HEADER_BYTES: usize = 56;
+
+struct BinaryMultiLoraRequest {
+    tensors: rustrain_ipc::TensorSlabRef,
+    payload: Vec<u8>,
+    batch_size: usize,
+    seq_len: usize,
+    n_total: i32,
+    lora_rank: i32,
+    adapter_ids: Vec<i64>,
+    expected_steps: Vec<u64>,
+}
+
+fn read_binary_u16(bytes: &[u8], offset: &mut usize) -> Result<u16, String> {
+    let end = offset
+        .checked_add(2)
+        .ok_or_else(|| "binary header offset overflowed".to_string())?;
+    let value = bytes
+        .get(*offset..end)
+        .ok_or_else(|| "binary multi-LoRA header is truncated".to_string())?;
+    *offset = end;
+    Ok(u16::from_le_bytes([value[0], value[1]]))
+}
+
+fn read_binary_u32(bytes: &[u8], offset: &mut usize) -> Result<u32, String> {
+    let end = offset
+        .checked_add(4)
+        .ok_or_else(|| "binary header offset overflowed".to_string())?;
+    let value = bytes
+        .get(*offset..end)
+        .ok_or_else(|| "binary multi-LoRA header is truncated".to_string())?;
+    *offset = end;
+    Ok(u32::from_le_bytes(value.try_into().unwrap()))
+}
+
+fn read_binary_i32(bytes: &[u8], offset: &mut usize) -> Result<i32, String> {
+    Ok(read_binary_u32(bytes, offset)? as i32)
+}
+
+fn read_binary_u64(bytes: &[u8], offset: &mut usize) -> Result<u64, String> {
+    let end = offset
+        .checked_add(8)
+        .ok_or_else(|| "binary header offset overflowed".to_string())?;
+    let value = bytes
+        .get(*offset..end)
+        .ok_or_else(|| "binary multi-LoRA header is truncated".to_string())?;
+    *offset = end;
+    Ok(u64::from_le_bytes(value.try_into().unwrap()))
+}
+
+fn read_binary_i64_vec(
+    bytes: &[u8],
+    offset: &mut usize,
+    count: usize,
+    label: &str,
+) -> Result<Vec<i64>, String> {
+    let byte_count = count
+        .checked_mul(std::mem::size_of::<i64>())
+        .ok_or_else(|| format!("binary {label} count overflowed"))?;
+    let end = offset
+        .checked_add(byte_count)
+        .ok_or_else(|| format!("binary {label} offset overflowed"))?;
+    let value = bytes
+        .get(*offset..end)
+        .ok_or_else(|| format!("binary {label} section is truncated"))?;
+    *offset = end;
+    value
+        .chunks_exact(8)
+        .map(|chunk| Ok(i64::from_le_bytes(chunk.try_into().unwrap())))
+        .collect()
+}
+
+fn parse_binary_multi_lora_request(bytes: &[u8]) -> Result<BinaryMultiLoraRequest, String> {
+    if bytes.len() < MULTI_LORA_BINARY_HEADER_BYTES {
+        return Err("binary multi-LoRA request is shorter than its header".to_string());
+    }
+    if bytes[..4] != MULTI_LORA_BINARY_MAGIC {
+        return Err("binary multi-LoRA request has an invalid magic".to_string());
+    }
+    let mut offset = 4;
+    let version = read_binary_u16(bytes, &mut offset)?;
+    let flags = read_binary_u16(bytes, &mut offset)?;
+    if version != MULTI_LORA_BINARY_VERSION {
+        return Err(format!(
+            "unsupported binary multi-LoRA version {version}, expected {MULTI_LORA_BINARY_VERSION}"
+        ));
+    }
+    if flags != 0 {
+        return Err("binary multi-LoRA request contains unsupported flags".to_string());
+    }
+    let batch_size = usize::try_from(read_binary_u32(bytes, &mut offset)?)
+        .map_err(|_| "binary batch size does not fit usize".to_string())?;
+    let seq_len = usize::try_from(read_binary_u32(bytes, &mut offset)?)
+        .map_err(|_| "binary sequence length does not fit usize".to_string())?;
+    let n_total = read_binary_i32(bytes, &mut offset)?;
+    let lora_rank = read_binary_i32(bytes, &mut offset)?;
+    let adapter_count = usize::try_from(read_binary_u32(bytes, &mut offset)?)
+        .map_err(|_| "binary adapter count does not fit usize".to_string())?;
+    let expected_count = usize::try_from(read_binary_u32(bytes, &mut offset)?)
+        .map_err(|_| "binary expected-step count does not fit usize".to_string())?;
+    let input_bytes = usize::try_from(read_binary_u64(bytes, &mut offset)?)
+        .map_err(|_| "binary input length does not fit usize".to_string())?;
+    let target_bytes = usize::try_from(read_binary_u64(bytes, &mut offset)?)
+        .map_err(|_| "binary target length does not fit usize".to_string())?;
+    let attention_bytes = usize::try_from(read_binary_u64(bytes, &mut offset)?)
+        .map_err(|_| "binary attention length does not fit usize".to_string())?;
+    if offset != MULTI_LORA_BINARY_HEADER_BYTES {
+        return Err("binary multi-LoRA header size mismatch".to_string());
+    }
+    if batch_size == 0 || seq_len <= 1 || n_total <= 0 || lora_rank <= 0 {
+        return Err("binary multi-LoRA dimensions and counts must be positive".to_string());
+    }
+    let n_total_usize = usize::try_from(n_total)
+        .map_err(|_| "binary adapter count does not fit usize".to_string())?;
+    if adapter_count != n_total_usize {
+        return Err(format!(
+            "binary adapter count {adapter_count} does not match n_total {n_total}"
+        ));
+    }
+    if expected_count != 0 && expected_count != n_total_usize {
+        return Err("binary expected-step count must be zero or n_total".to_string());
+    }
+    let elements = batch_size
+        .checked_mul(seq_len)
+        .ok_or_else(|| "binary tensor element count overflowed".to_string())?;
+    let expected_tensor_bytes = elements
+        .checked_mul(std::mem::size_of::<i64>())
+        .ok_or_else(|| "binary tensor byte count overflowed".to_string())?;
+    if input_bytes != expected_tensor_bytes
+        || target_bytes != expected_tensor_bytes
+        || attention_bytes != expected_tensor_bytes
+    {
+        return Err(format!(
+            "binary tensor lengths must all equal {expected_tensor_bytes}"
+        ));
+    }
+    let adapter_ids = read_binary_i64_vec(bytes, &mut offset, n_total_usize, "adapter IDs")?;
+    let expected_steps = if expected_count == 0 {
+        Vec::new()
+    } else {
+        let values = read_binary_i64_vec(bytes, &mut offset, expected_count, "expected steps")?;
+        values
+            .into_iter()
+            .map(|value| {
+                u64::try_from(value)
+                    .map_err(|_| "binary expected steps must be non-negative".to_string())
+            })
+            .collect::<Result<Vec<_>, _>>()?
+    };
+    let tensor_sections = [input_bytes, target_bytes, attention_bytes];
+    let mut payload = Vec::with_capacity(
+        tensor_sections
+            .iter()
+            .try_fold(0usize, |total, size| total.checked_add(*size))
+            .and_then(|total| total.checked_add(2 * (rustrain_ipc::TENSOR_SPAN_ALIGNMENT - 1)))
+            .ok_or_else(|| "binary tensor slab capacity overflowed".to_string())?,
+    );
+    let mut spans = Vec::with_capacity(3);
+    for section_size in tensor_sections {
+        let aligned = payload
+            .len()
+            .checked_add(rustrain_ipc::TENSOR_SPAN_ALIGNMENT - 1)
+            .map(|value| value & !(rustrain_ipc::TENSOR_SPAN_ALIGNMENT - 1))
+            .ok_or_else(|| "binary tensor slab alignment overflowed".to_string())?;
+        payload.resize(aligned, 0);
+        let start = offset;
+        let end = start
+            .checked_add(section_size)
+            .ok_or_else(|| "binary tensor section offset overflowed".to_string())?;
+        payload.extend_from_slice(
+            bytes
+                .get(start..end)
+                .ok_or_else(|| "binary tensor section is truncated".to_string())?,
+        );
+        offset = end;
+        spans.push(rustrain_ipc::TensorSpan {
+            offset_bytes: u64::try_from(aligned)
+                .map_err(|_| "binary tensor span offset exceeds u64".to_string())?,
+            len_bytes: u64::try_from(section_size)
+                .map_err(|_| "binary tensor span length exceeds u64".to_string())?,
+        });
+    }
+    if offset != bytes.len() {
+        return Err("binary multi-LoRA request has trailing bytes".to_string());
+    }
+    let tensors = rustrain_ipc::TensorSlabRef {
+        input_ids: spans[0],
+        target_mask: spans[1],
+        attention_mask: spans[2],
+        batch_size,
+        seq_len,
+    };
+    tensors.validate(payload.len())?;
+    Ok(BinaryMultiLoraRequest {
+        tensors,
+        payload,
+        batch_size,
+        seq_len,
+        n_total,
+        lora_rank,
+        adapter_ids,
+        expected_steps,
+    })
 }
 #[derive(Serialize)]
 struct TrainStepResponse {
@@ -252,16 +979,402 @@ struct TrainStepResponse {
     step: u64,
 }
 
-/// Decode a base64-encoded int64 tensor to Vec<i64> (for EP IPC).
-fn decode_int64_vec(t: &TensorHttp) -> Result<Vec<i64>, String> {
+#[derive(Serialize)]
+struct TrainMultiLoraResponse {
+    capability_version: u32,
+    capabilities: &'static [&'static str],
+    loss: f64,
+    step: u64,
+    loss_scope: &'static str,
+    coalesced_requests: usize,
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    adapter_losses: Vec<rustrain_ipc::command::AdapterLoss>,
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    adapter_steps: Vec<rustrain_ipc::command::AdapterStep>,
+}
+
+fn multi_lora_capability_v1(headers: &HeaderMap) -> bool {
+    headers
+        .get(MULTI_LORA_CAPABILITY_HEADER)
+        .and_then(|value| value.to_str().ok())
+        .is_some_and(|value| {
+            value
+                .split(',')
+                .map(str::trim)
+                .any(|token| token.eq_ignore_ascii_case(MULTI_LORA_CAPABILITY_V1))
+        })
+}
+
+/// Decode the little-endian wire representation without materializing host values.
+fn decode_int64_bytes(t: &TensorHttp) -> Result<Vec<u8>, String> {
     use base64::{engine::general_purpose, Engine};
+    if t.dtype != "int64" {
+        return Err(format!("EP tensor dtype must be int64, got {}", t.dtype));
+    }
     let bytes = general_purpose::STANDARD
         .decode(&t.data)
         .map_err(|e| format!("base64 decode: {e}"))?;
-    Ok(bytes
-        .chunks_exact(8)
-        .map(|c| i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]))
-        .collect())
+    let expected = t
+        .shape
+        .iter()
+        .try_fold(1usize, |acc, dim| {
+            usize::try_from(*dim)
+                .ok()
+                .and_then(|dim| acc.checked_mul(dim))
+        })
+        .ok_or_else(|| format!("invalid tensor shape {:?}", t.shape))?;
+    let expected_bytes = expected
+        .checked_mul(std::mem::size_of::<i64>())
+        .ok_or_else(|| format!("tensor shape {:?} byte size overflows usize", t.shape))?;
+    if expected_bytes != bytes.len() {
+        return Err(format!(
+            "tensor shape {:?} expects {} int64 bytes, got {}",
+            t.shape,
+            expected_bytes,
+            bytes.len()
+        ));
+    }
+    Ok(bytes)
+}
+
+fn decode_int64_values(t: &TensorHttp) -> Result<Vec<i64>, String> {
+    let bytes = decode_int64_bytes(t)?;
+    bytes
+        .chunks_exact(std::mem::size_of::<i64>())
+        .map(|chunk| {
+            let bytes: [u8; 8] = chunk
+                .try_into()
+                .map_err(|_| "invalid int64 byte width".to_string())?;
+            Ok(i64::from_le_bytes(bytes))
+        })
+        .collect()
+}
+
+fn pack_tensor_slab(
+    input_ids: &TensorHttp,
+    target_mask: &TensorHttp,
+    attention_mask: &TensorHttp,
+    batch_size: usize,
+    seq_len: usize,
+) -> Result<(rustrain_ipc::TensorSlabRef, Vec<u8>), String> {
+    let decoded = [
+        decode_int64_bytes(input_ids)?,
+        decode_int64_bytes(target_mask)?,
+        decode_int64_bytes(attention_mask)?,
+    ];
+    let total = decoded
+        .iter()
+        .try_fold(0usize, |sum, bytes| sum.checked_add(bytes.len()))
+        .and_then(|sum| sum.checked_add(2 * (rustrain_ipc::TENSOR_SPAN_ALIGNMENT - 1)))
+        .ok_or_else(|| "tensor slab payload size overflows usize".to_string())?;
+    let mut payload = Vec::with_capacity(total);
+    let mut spans = Vec::with_capacity(3);
+    for bytes in decoded {
+        let aligned = payload
+            .len()
+            .checked_add(rustrain_ipc::TENSOR_SPAN_ALIGNMENT - 1)
+            .map(|value| value & !(rustrain_ipc::TENSOR_SPAN_ALIGNMENT - 1))
+            .ok_or_else(|| "tensor slab alignment overflows usize".to_string())?;
+        payload.resize(aligned, 0);
+        let offset_bytes =
+            u64::try_from(aligned).map_err(|_| "tensor slab offset exceeds u64".to_string())?;
+        let len_bytes = u64::try_from(bytes.len())
+            .map_err(|_| "tensor slab tensor length exceeds u64".to_string())?;
+        payload.extend_from_slice(&bytes);
+        spans.push(rustrain_ipc::TensorSpan {
+            offset_bytes,
+            len_bytes,
+        });
+    }
+    Ok((
+        rustrain_ipc::TensorSlabRef {
+            input_ids: spans[0],
+            target_mask: spans[1],
+            attention_mask: spans[2],
+            batch_size,
+            seq_len,
+        },
+        payload,
+    ))
+}
+
+fn multi_lora_source_count(batch_size: usize, n_total: usize) -> Result<usize, String> {
+    if n_total == 0 {
+        return Err("multi-LoRA adapter count must be positive".to_string());
+    }
+    if batch_size == 1 {
+        return Ok(1);
+    }
+    if batch_size % n_total != 0 {
+        return Err(format!(
+            "multi-LoRA batch_size={batch_size} must be 1 or a multiple of n_total={n_total}"
+        ));
+    }
+    Ok(batch_size / n_total)
+}
+
+fn normalized_multi_lora_payload_bytes(
+    seq_len: usize,
+    n_total: usize,
+    source_count: usize,
+) -> Result<usize, String> {
+    n_total
+        .checked_mul(source_count)
+        .and_then(|rows| rows.checked_mul(seq_len))
+        .and_then(|elements| elements.checked_mul(std::mem::size_of::<i64>()))
+        .and_then(|tensor_bytes| tensor_bytes.checked_mul(3))
+        .and_then(|bytes| bytes.checked_add(2 * (rustrain_ipc::TENSOR_SPAN_ALIGNMENT - 1)))
+        .ok_or_else(|| "normalized multi-LoRA payload size overflowed usize".to_string())
+}
+
+type MultiLoraWindowBuild = (
+    rustrain_ipc::EpCommand,
+    Vec<u8>,
+    Vec<MultiLoraResponseTarget>,
+);
+
+fn build_multi_lora_window(
+    window: MultiLoraWindow,
+) -> Result<MultiLoraWindowBuild, (String, Vec<MultiLoraResponseTarget>)> {
+    let build_result = build_multi_lora_window_payload(&window);
+    let mut adapter_start = 0usize;
+    let responses = window
+        .requests
+        .into_iter()
+        .map(|request| {
+            let adapter_end = adapter_start + request.n_total;
+            let target = MultiLoraResponseTarget {
+                response: request.response,
+                adapter_range: adapter_start..adapter_end,
+            };
+            adapter_start = adapter_end;
+            target
+        })
+        .collect::<Vec<_>>();
+    match build_result {
+        Ok((tensors, payload, adapter_ids, lora_rank, expected_steps)) => {
+            let n_total = match i32::try_from(adapter_ids.len()) {
+                Ok(n_total) => n_total,
+                Err(_) => {
+                    return Err(("coalesced adapter count exceeds i32".to_string(), responses));
+                }
+            };
+            Ok((
+                rustrain_ipc::EpCommand::TrainMultiLoraSlab {
+                    session_id: window.session_id,
+                    tensors,
+                    n_total,
+                    lora_rank,
+                    adapter_ids,
+                    expected_steps,
+                },
+                payload,
+                responses,
+            ))
+        }
+        Err(error) => Err((error, responses)),
+    }
+}
+
+fn build_multi_lora_window_payload(
+    window: &MultiLoraWindow,
+) -> Result<
+    (
+        rustrain_ipc::TensorSlabRef,
+        Vec<u8>,
+        Vec<i64>,
+        i32,
+        Vec<u64>,
+    ),
+    String,
+> {
+    if window.requests.is_empty() {
+        return Err("coalesced multi-LoRA window is empty".to_string());
+    }
+    let total_adapters = window
+        .requests
+        .iter()
+        .try_fold(0usize, |total, request| total.checked_add(request.n_total))
+        .ok_or_else(|| "coalesced adapter count overflowed usize".to_string())?;
+    let batch_size = total_adapters
+        .checked_mul(window.source_count)
+        .ok_or_else(|| "coalesced batch size overflowed usize".to_string())?;
+    let row_bytes = window
+        .seq_len
+        .checked_mul(std::mem::size_of::<i64>())
+        .ok_or_else(|| "coalesced row byte count overflowed usize".to_string())?;
+    let tensor_bytes = batch_size
+        .checked_mul(row_bytes)
+        .ok_or_else(|| "coalesced tensor byte count overflowed usize".to_string())?;
+    let payload_capacity = tensor_bytes
+        .checked_mul(3)
+        .and_then(|bytes| bytes.checked_add(2 * (rustrain_ipc::TENSOR_SPAN_ALIGNMENT - 1)))
+        .ok_or_else(|| "coalesced payload capacity overflowed usize".to_string())?;
+    let mut payload = Vec::with_capacity(payload_capacity);
+    let mut spans = Vec::with_capacity(3);
+
+    for tensor_index in 0..3 {
+        let aligned = payload
+            .len()
+            .checked_add(rustrain_ipc::TENSOR_SPAN_ALIGNMENT - 1)
+            .map(|value| value & !(rustrain_ipc::TENSOR_SPAN_ALIGNMENT - 1))
+            .ok_or_else(|| "coalesced tensor alignment overflowed usize".to_string())?;
+        payload.resize(aligned, 0);
+        let offset_bytes = u64::try_from(aligned)
+            .map_err(|_| "coalesced tensor offset exceeds u64".to_string())?;
+        for source_index in 0..window.source_count {
+            for request in &window.requests {
+                request.tensors.validate(request.payload.len())?;
+                if request.source_count != window.source_count
+                    || request.tensors.seq_len != window.seq_len
+                {
+                    return Err("coalesced request layout changed after admission".to_string());
+                }
+                let span = request.tensors.spans()[tensor_index];
+                let tensor_start = usize::try_from(span.offset_bytes)
+                    .map_err(|_| "request tensor offset exceeds usize".to_string())?;
+                for adapter_index in 0..request.n_total {
+                    let row_index = if request.tensors.batch_size == 1 {
+                        0
+                    } else {
+                        source_index
+                            .checked_mul(request.n_total)
+                            .and_then(|row| row.checked_add(adapter_index))
+                            .ok_or_else(|| "request row index overflowed usize".to_string())?
+                    };
+                    let start = tensor_start
+                        .checked_add(
+                            row_index
+                                .checked_mul(row_bytes)
+                                .ok_or_else(|| "request row offset overflowed usize".to_string())?,
+                        )
+                        .ok_or_else(|| "request tensor range overflowed usize".to_string())?;
+                    let end = start
+                        .checked_add(row_bytes)
+                        .ok_or_else(|| "request tensor range overflowed usize".to_string())?;
+                    let row = request
+                        .payload
+                        .get(start..end)
+                        .ok_or_else(|| "request tensor row exceeds payload".to_string())?;
+                    payload.extend_from_slice(row);
+                }
+            }
+        }
+        spans.push(rustrain_ipc::TensorSpan {
+            offset_bytes,
+            len_bytes: u64::try_from(tensor_bytes)
+                .map_err(|_| "coalesced tensor length exceeds u64".to_string())?,
+        });
+    }
+
+    let adapter_ids = window
+        .requests
+        .iter()
+        .flat_map(|request| request.adapter_ids.iter().copied())
+        .collect::<Vec<_>>();
+    if adapter_ids.len() != total_adapters && !adapter_ids.is_empty() {
+        return Err("coalesced adapter ID count does not match batch geometry".to_string());
+    }
+    let expected_steps = window
+        .requests
+        .iter()
+        .flat_map(|request| request.expected_steps.iter().copied())
+        .collect::<Vec<_>>();
+    if !expected_steps.is_empty() && expected_steps.len() != total_adapters {
+        return Err("expected step count does not match coalesced adapter count".to_string());
+    }
+    let lora_rank = window
+        .requests
+        .iter()
+        .map(|request| request.lora_rank)
+        .max()
+        .unwrap_or(0);
+    let tensors = rustrain_ipc::TensorSlabRef {
+        input_ids: spans[0],
+        target_mask: spans[1],
+        attention_mask: spans[2],
+        batch_size,
+        seq_len: window.seq_len,
+    };
+    tensors.validate(payload.len())?;
+    Ok((tensors, payload, adapter_ids, lora_rank, expected_steps))
+}
+
+fn validate_train_http_shapes(
+    input_ids: &TensorHttp,
+    target_mask: &TensorHttp,
+    attention_mask: &TensorHttp,
+) -> Result<(usize, usize), String> {
+    if input_ids.shape.len() != 1 && input_ids.shape.len() != 2 {
+        return Err(format!(
+            "input_ids must have shape [seq] or [batch, seq], got {:?}",
+            input_ids.shape
+        ));
+    }
+    if target_mask.shape != input_ids.shape || attention_mask.shape != input_ids.shape {
+        return Err(format!(
+            "masks must have the same shape as input_ids: input={:?} target={:?} attention={:?}",
+            input_ids.shape, target_mask.shape, attention_mask.shape
+        ));
+    }
+    let seq_len = *input_ids
+        .shape
+        .last()
+        .ok_or_else(|| "input_ids shape is empty".to_string())?;
+    if seq_len <= 0 {
+        return Err(format!("sequence length must be positive, got {seq_len}"));
+    }
+    let batch_size = if input_ids.shape.len() == 2 {
+        input_ids.shape[0]
+    } else {
+        1
+    };
+    if batch_size <= 0 {
+        return Err(format!("batch size must be positive, got {batch_size}"));
+    }
+    Ok((
+        usize::try_from(batch_size).map_err(|_| format!("invalid batch size {batch_size}"))?,
+        usize::try_from(seq_len).map_err(|_| format!("invalid sequence length {seq_len}"))?,
+    ))
+}
+
+fn validate_multi_lora_http_shapes(
+    input_ids: &TensorHttp,
+    target_mask: &TensorHttp,
+    attention_mask: &TensorHttp,
+    n_total: i32,
+) -> Result<(usize, usize), String> {
+    if n_total <= 0 {
+        return Err(format!("n_total must be positive, got {n_total}"));
+    }
+    let (batch_size, seq_len) = validate_train_http_shapes(input_ids, target_mask, attention_mask)?;
+    let n_total = n_total as usize;
+    if batch_size != 1 && batch_size % n_total != 0 {
+        return Err(format!(
+            "multi-LoRA batch must be 1 or a positive multiple of n_total={n_total}, got {batch_size}"
+        ));
+    }
+    Ok((batch_size, seq_len))
+}
+
+fn validate_selected_eval_http_shapes(
+    input_ids: &TensorHttp,
+    target_mask: &TensorHttp,
+    attention_mask: &TensorHttp,
+    adapter_ids: &[i64],
+) -> Result<(usize, usize), String> {
+    if adapter_ids.is_empty() || adapter_ids.iter().any(|id| *id <= 0) {
+        return Err("adapter_ids must contain positive IDs".to_string());
+    }
+    let (batch_size, seq_len) = validate_train_http_shapes(input_ids, target_mask, attention_mask)?;
+    if batch_size != 1 && batch_size != adapter_ids.len() {
+        return Err(format!(
+            "selected eval batch_size must be 1 or adapter count {}, got {}",
+            adapter_ids.len(),
+            batch_size
+        ));
+    }
+    Ok((batch_size, seq_len))
 }
 
 fn decode_tensor(t: &TensorHttp) -> Result<tch::Tensor, String> {
@@ -293,8 +1406,13 @@ fn decode_tensor(t: &TensorHttp) -> Result<tch::Tensor, String> {
         }
         _ => return Err("only int64 and float32 supported via HTTP".into()),
     };
-    let local_rank = std::env::var("LOCAL_RANK").ok().and_then(|s| s.parse::<usize>().ok()).unwrap_or(0);
-    Ok(tensor.reshape(&t.shape).to_device(tch::Device::Cuda(local_rank)))
+    let local_rank = std::env::var("LOCAL_RANK")
+        .ok()
+        .and_then(|s| s.parse::<usize>().ok())
+        .unwrap_or(0);
+    Ok(tensor
+        .reshape(&t.shape)
+        .to_device(tch::Device::Cuda(local_rank)))
 }
 
 async fn eval_step(
@@ -320,6 +1438,52 @@ async fn eval_step(
         .map_err(|e| err_resp(&e.to_string()))?;
     Ok(Json(EvalStepResponse { loss: result.loss }))
 }
+
+async fn eval_multi_lora(
+    State(state): State<Arc<AppState>>,
+    Path(id): Path<String>,
+    Json(req): Json<EvalMultiLoraHttp>,
+) -> Result<Json<TrainMultiLoraEvalResponse>, (StatusCode, Json<ErrorResponse>)> {
+    let (batch_size, seq_len) = validate_selected_eval_http_shapes(
+        &req.input_ids,
+        &req.target_mask,
+        &req.attention_mask,
+        &req.adapter_ids,
+    )
+    .map_err(|e| err_resp(&e))?;
+    let input_ids = decode_int64_values(&req.input_ids).map_err(|e| err_resp(&e))?;
+    let target_mask = decode_int64_values(&req.target_mask).map_err(|e| err_resp(&e))?;
+    let attention_mask = decode_int64_values(&req.attention_mask).map_err(|e| err_resp(&e))?;
+    let session = state
+        .manager
+        .get_session(&id)
+        .await
+        .ok_or_else(|| err_resp("session not found"))?;
+    let s = session.lock().await;
+    let output = s
+        .eval_multi_lora_host_i64(
+            &input_ids,
+            &target_mask,
+            &attention_mask,
+            batch_size,
+            seq_len,
+            &req.adapter_ids,
+        )
+        .map_err(|e| err_resp(&e.to_string()))?;
+    Ok(Json(TrainMultiLoraEvalResponse {
+        adapter_losses: output
+            .adapter_losses
+            .into_iter()
+            .map(|(adapter_id, loss)| rustrain_ipc::command::AdapterLoss { adapter_id, loss })
+            .collect(),
+    }))
+}
+
+#[derive(Serialize)]
+struct TrainMultiLoraEvalResponse {
+    adapter_losses: Vec<rustrain_ipc::command::AdapterLoss>,
+}
+
 #[derive(Serialize)]
 struct EvalStepResponse {
     loss: f64,
@@ -381,6 +1545,8 @@ async fn load_checkpoint(
 #[derive(Deserialize)]
 struct ExportHttp {
     path: String,
+    #[serde(default)]
+    adapter_id: Option<i64>,
 }
 #[derive(Serialize)]
 struct ExportResponse {
@@ -400,7 +1566,7 @@ async fn export_adapter(
         .ok_or_else(|| err_resp("session not found"))?;
     let s = session.lock().await;
     let count = s
-        .export_adapter(&req.path)
+        .export_adapter(&req.path, req.adapter_id)
         .map_err(|e| err_resp(&e.to_string()))?;
     Ok(Json(ExportResponse {
         path: req.path,
@@ -408,6 +1574,33 @@ async fn export_adapter(
     }))
 }
 
+#[derive(Deserialize)]
+struct ImportHttp {
+    path: String,
+}
+
+#[derive(Serialize)]
+struct ImportResponse {
+    adapter_id: i64,
+}
+
+async fn import_adapter(
+    State(state): State<Arc<AppState>>,
+    Path(id): Path<String>,
+    Json(req): Json<ImportHttp>,
+) -> Result<Json<ImportResponse>, (StatusCode, Json<ErrorResponse>)> {
+    let session = state
+        .manager
+        .get_session(&id)
+        .await
+        .ok_or_else(|| err_resp("session not found"))?;
+    let mut s = session.lock().await;
+    let adapter_id = s
+        .import_adapter(&req.path)
+        .map_err(|e| err_resp(&e.to_string()))?;
+    Ok(Json(ImportResponse { adapter_id }))
+}
+
 async fn get_status(
     State(state): State<Arc<AppState>>,
     Path(id): Path<String>,
@@ -433,6 +1626,14 @@ struct AddLoRAHttp {
     alpha: f64,
     target_layers: Vec<i64>,
     target_modules: String,
+    #[serde(default)]
+    optimizer_lr: Option<f64>,
+    #[serde(default)]
+    optimizer_beta1: Option<f64>,
+    #[serde(default)]
+    optimizer_beta2: Option<f64>,
+    #[serde(default)]
+    optimizer_eps: Option<f64>,
 }
 #[derive(Serialize)]
 struct AddLoRAResponse {
@@ -456,6 +1657,10 @@ async fn add_lora(
             alpha: req.alpha,
             target_layers: req.target_layers,
             target_modules: req.target_modules,
+            optimizer_lr: req.optimizer_lr,
+            optimizer_beta1: req.optimizer_beta1,
+            optimizer_beta2: req.optimizer_beta2,
+            optimizer_eps: req.optimizer_eps,
         })
         .map_err(|e| err_resp(&e.to_string()))?;
     Ok(Json(AddLoRAResponse { adapter_id }))
@@ -522,13 +1727,15 @@ async fn stream_metrics(
     drop(s);
 
     let stream = stream::iter(metrics.into_iter().map(|m| {
-        Ok(Event::default().json_data(StepMetricJson {
-            step: m.step,
-            loss: m.loss,
-            lr: m.lr,
-            mem_gb: m.mem_gb,
-            timestamp_unix: m.timestamp_unix,
-        }).unwrap_or_default())
+        Ok(Event::default()
+            .json_data(StepMetricJson {
+                step: m.step,
+                loss: m.loss,
+                lr: m.lr,
+                mem_gb: m.mem_gb,
+                timestamp_unix: m.timestamp_unix,
+            })
+            .unwrap_or_default())
     }));
 
     Ok(Sse::new(stream).keep_alive(KeepAlive::default()))
@@ -548,15 +1755,56 @@ struct StepMetricJson {
 // ──────────────────────────────────────────────────────────────────────
 
 pub fn ep_router(state: Arc<EpAppState>) -> Router {
+    let batch_config = MultiLoraBatchConfig::from_env();
+    let multi_lora_batcher = MultiLoraBatcher::new(batch_config);
+    let state = Arc::new(EpRouterState {
+        coordinator: Arc::clone(&state.coordinator),
+        world_size: state.world_size,
+        dispatcher: EpDispatchScheduler::new(configured_queue_capacity()),
+        dispatch_submission: Mutex::new(()),
+        multi_lora_batcher,
+    });
+    let tensor_routes = Router::new()
+        .route("/v1/sessions/{id}/train_step", post(ep_train_step))
+        .route("/v1/sessions/{id}/eval_step", post(ep_eval_step))
+        .route(
+            "/v1/sessions/{id}/eval_multi_lora",
+            post(ep_eval_multi_lora),
+        )
+        .route_layer(middleware::from_fn_with_state(
+            Arc::new(Semaphore::new(1)),
+            ep_tensor_admission,
+        ));
+    let multi_lora_route = Router::new()
+        .route("/v1/sessions/{id}/train_multi", post(ep_train_multi_lora))
+        .route(
+            "/v1/sessions/{id}/train_multi_binary",
+            post(ep_train_multi_lora_binary),
+        )
+        .route_layer(middleware::from_fn_with_state(
+            Arc::new(Semaphore::new(batch_config.max_requests)),
+            ep_tensor_admission,
+        ));
+
     Router::new()
+        .merge(tensor_routes)
+        .merge(multi_lora_route)
         .route("/v1/sessions", post(ep_create_session))
-        .route("/v1/sessions/{id}", axum::routing::delete(ep_delete_session))
+        .route(
+            "/v1/sessions/{id}",
+            axum::routing::delete(ep_delete_session),
+        )
         .route("/v1/sessions/{id}/load_model", post(ep_load_model))
         .route("/v1/sessions/{id}/load_dataset", post(ep_load_dataset))
         .route("/v1/sessions/{id}/init_lora", post(ep_init_lora))
-        .route("/v1/sessions/{id}/train_step", post(ep_train_step))
-        .route("/v1/sessions/{id}/train_multi", post(ep_train_multi_lora))
-        .route("/v1/sessions/{id}/eval_step", post(ep_eval_step))
+        .route(
+            "/v1/sessions/{id}/save_checkpoint",
+            post(ep_save_checkpoint),
+        )
+        .route(
+            "/v1/sessions/{id}/load_checkpoint",
+            post(ep_load_checkpoint),
+        )
         .route("/v1/sessions/{id}/add_lora", post(ep_add_lora))
         .route("/v1/sessions/{id}/batch_add_lora", post(ep_batch_add_lora))
         .route("/v1/sessions/{id}/remove_lora", post(ep_remove_lora))
@@ -564,31 +1812,69 @@ pub fn ep_router(state: Arc<EpAppState>) -> Router {
         .route("/v1/sessions/{id}/export_adapter", post(ep_export_adapter))
         .route("/v1/sessions/{id}/status", get(ep_get_status))
         .route("/v1/health", get(ep_health))
+        .layer(DefaultBodyLimit::max(48 * 1024 * 1024))
         .with_state(state)
 }
 
-async fn ep_health() -> Json<serde_json::Value> {
-    Json(serde_json::json!({"status": "ok", "mode": "ep"}))
+async fn ep_tensor_admission(
+    State(gate): State<Arc<Semaphore>>,
+    request: Request,
+    next: Next,
+) -> Response {
+    // The IPC coordinator serializes dispatches, so decoding more than one tensor request only
+    // increases peak memory without creating training concurrency.
+    match gate.try_acquire_owned() {
+        Ok(_permit) => next.run(request).await,
+        Err(_) => (
+            StatusCode::TOO_MANY_REQUESTS,
+            Json(ErrorResponse {
+                error: "EP tensor admission capacity is exhausted".to_string(),
+            }),
+        )
+            .into_response(),
+    }
+}
+
+async fn ep_health(
+    State(state): State<Arc<EpRouterState>>,
+) -> (StatusCode, Json<serde_json::Value>) {
+    if state.coordinator.is_healthy() {
+        (
+            StatusCode::OK,
+            Json(serde_json::json!({"status": "ok", "mode": "ep"})),
+        )
+    } else {
+        (
+            StatusCode::SERVICE_UNAVAILABLE,
+            Json(serde_json::json!({"status": "error", "mode": "ep"})),
+        )
+    }
 }
 
 async fn ep_create_session(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Json(req): Json<CreateSessionRequest>,
 ) -> Result<Json<CreateSessionResponse>, (StatusCode, Json<ErrorResponse>)> {
-    let cmd = rustrain_ipc::EpCommand::CreateSession { session_id: req.session_id.clone() };
-    match state.coordinator.dispatch(&cmd) {
-        rustrain_ipc::EpResult::Ok => Ok(Json(CreateSessionResponse { session_id: req.session_id })),
+    let cmd = rustrain_ipc::EpCommand::CreateSession {
+        session_id: req.session_id.clone(),
+    };
+    match dispatch_ep(&state, cmd).await? {
+        rustrain_ipc::EpResult::Ok => Ok(Json(CreateSessionResponse {
+            session_id: req.session_id,
+        })),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
     }
 }
 
 async fn ep_delete_session(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
 ) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
-    let cmd = rustrain_ipc::EpCommand::DeleteSession { session_id: id.clone() };
-    match state.coordinator.dispatch(&cmd) {
+    let cmd = rustrain_ipc::EpCommand::DeleteSession {
+        session_id: id.clone(),
+    };
+    match dispatch_ep(&state, cmd).await? {
         rustrain_ipc::EpResult::Ok => Ok(Json(serde_json::json!({"deleted": id}))),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
@@ -596,7 +1882,7 @@ async fn ep_delete_session(
 }
 
 async fn ep_load_model(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
     Json(req): Json<LoadModelHttp>,
 ) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
@@ -605,7 +1891,7 @@ async fn ep_load_model(
         model_path: req.model_path,
         config_toml: req.config_toml,
     };
-    match state.coordinator.dispatch(&cmd) {
+    match dispatch_ep(&state, cmd).await? {
         rustrain_ipc::EpResult::Ok => Ok(Json(serde_json::json!({"loaded": true}))),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
@@ -613,7 +1899,7 @@ async fn ep_load_model(
 }
 
 async fn ep_load_dataset(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
     Json(req): Json<LoadDatasetHttp>,
 ) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
@@ -622,7 +1908,7 @@ async fn ep_load_dataset(
         jsonl_path: req.jsonl_path,
         seq_len: req.seq_len,
     };
-    match state.coordinator.dispatch(&cmd) {
+    match dispatch_ep(&state, cmd).await? {
         rustrain_ipc::EpResult::Count(n) => Ok(Json(serde_json::json!({"samples": n}))),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
@@ -630,7 +1916,7 @@ async fn ep_load_dataset(
 }
 
 async fn ep_init_lora(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
     Json(req): Json<InitLoRAHttp>,
 ) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
@@ -645,7 +1931,7 @@ async fn ep_init_lora(
         beta2: req.beta2,
         eps: req.eps,
     };
-    match state.coordinator.dispatch(&cmd) {
+    match dispatch_ep(&state, cmd).await? {
         rustrain_ipc::EpResult::Count(n) => Ok(Json(serde_json::json!({"lora_count": n}))),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
@@ -653,81 +1939,1055 @@ async fn ep_init_lora(
 }
 
 async fn ep_train_step(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
     Json(req): Json<TrainStepHttp>,
 ) -> Result<Json<TrainStepResponse>, (StatusCode, Json<ErrorResponse>)> {
-    let input_ids = decode_int64_vec(&req.input_ids).map_err(|e| err_resp(&e))?;
-    let target_mask = decode_int64_vec(&req.target_mask).map_err(|e| err_resp(&e))?;
-    let attention_mask = decode_int64_vec(&req.attention_mask).map_err(|e| err_resp(&e))?;
-    let seq_len = input_ids.len();
+    let (batch_size, seq_len) =
+        validate_train_http_shapes(&req.input_ids, &req.target_mask, &req.attention_mask)
+            .map_err(|e| err_resp(&e))?;
+    let (tensors, payload) = pack_tensor_slab(
+        &req.input_ids,
+        &req.target_mask,
+        &req.attention_mask,
+        batch_size,
+        seq_len,
+    )
+    .map_err(|e| err_resp(&e))?;
 
-    let cmd = rustrain_ipc::EpCommand::TrainStep {
+    let cmd = rustrain_ipc::EpCommand::TrainStepSlab {
         session_id: id,
-        input_ids,
-        target_mask,
-        attention_mask,
-        seq_len,
+        tensors,
     };
-    match state.coordinator.dispatch(&cmd) {
-        rustrain_ipc::EpResult::Loss(loss) => Ok(Json(TrainStepResponse { loss, step: 0 })),
+    match dispatch_slab(&state, cmd, payload).await? {
+        rustrain_ipc::EpResult::Train { loss, step } => Ok(Json(TrainStepResponse { loss, step })),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
     }
 }
 
 async fn ep_eval_step(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
     Json(req): Json<TrainStepHttp>,
 ) -> Result<Json<EvalStepResponse>, (StatusCode, Json<ErrorResponse>)> {
-    let input_ids = decode_int64_vec(&req.input_ids).map_err(|e| err_resp(&e))?;
-    let target_mask = decode_int64_vec(&req.target_mask).map_err(|e| err_resp(&e))?;
-    let attention_mask = decode_int64_vec(&req.attention_mask).map_err(|e| err_resp(&e))?;
-    let seq_len = input_ids.len();
+    let (batch_size, seq_len) =
+        validate_train_http_shapes(&req.input_ids, &req.target_mask, &req.attention_mask)
+            .map_err(|e| err_resp(&e))?;
+    let (tensors, payload) = pack_tensor_slab(
+        &req.input_ids,
+        &req.target_mask,
+        &req.attention_mask,
+        batch_size,
+        seq_len,
+    )
+    .map_err(|e| err_resp(&e))?;
 
-    let cmd = rustrain_ipc::EpCommand::EvalStep {
+    let cmd = rustrain_ipc::EpCommand::EvalStepSlab {
         session_id: id,
-        input_ids,
-        target_mask,
-        attention_mask,
-        seq_len,
+        tensors,
     };
-    match state.coordinator.dispatch(&cmd) {
+    match dispatch_slab(&state, cmd, payload).await? {
         rustrain_ipc::EpResult::Loss(loss) => Ok(Json(EvalStepResponse { loss })),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
     }
 }
 
+async fn ep_eval_multi_lora(
+    State(state): State<Arc<EpRouterState>>,
+    Path(id): Path<String>,
+    Json(req): Json<EvalMultiLoraHttp>,
+) -> Result<Json<TrainMultiLoraEvalResponse>, (StatusCode, Json<ErrorResponse>)> {
+    let (batch_size, seq_len) = validate_selected_eval_http_shapes(
+        &req.input_ids,
+        &req.target_mask,
+        &req.attention_mask,
+        &req.adapter_ids,
+    )
+    .map_err(|e| err_resp(&e))?;
+    let (tensors, payload) = pack_tensor_slab(
+        &req.input_ids,
+        &req.target_mask,
+        &req.attention_mask,
+        batch_size,
+        seq_len,
+    )
+    .map_err(|e| err_resp(&e))?;
+    let cmd = rustrain_ipc::EpCommand::EvalMultiLoraSlab {
+        session_id: id,
+        tensors,
+        adapter_ids: req.adapter_ids,
+    };
+    match dispatch_slab(&state, cmd, payload).await? {
+        rustrain_ipc::EpResult::MultiLoraEval { adapter_losses } => {
+            Ok(Json(TrainMultiLoraEvalResponse { adapter_losses }))
+        }
+        rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
+        _ => Err(err_resp("unexpected result")),
+    }
+}
+
 async fn ep_train_multi_lora(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
+    headers: HeaderMap,
     Json(req): Json<TrainMultiLoraHttp>,
-) -> Result<Json<TrainStepResponse>, (StatusCode, Json<ErrorResponse>)> {
-    let input_ids = decode_int64_vec(&req.input_ids).map_err(|e| err_resp(&e))?;
-    let target_mask = decode_int64_vec(&req.target_mask).map_err(|e| err_resp(&e))?;
-    let attention_mask = decode_int64_vec(&req.attention_mask).map_err(|e| err_resp(&e))?;
-    let seq_len = input_ids.len();
+) -> Result<Json<TrainMultiLoraResponse>, (StatusCode, Json<ErrorResponse>)> {
+    let (batch_size, seq_len) = validate_multi_lora_http_shapes(
+        &req.input_ids,
+        &req.target_mask,
+        &req.attention_mask,
+        req.n_total,
+    )
+    .map_err(|e| err_resp(&e))?;
+    crate::ep::validate_multi_lora_global_batch_size(batch_size, req.n_total, state.world_size)
+        .map_err(|error| err_resp(&error))?;
+    let (tensors, payload) = pack_tensor_slab(
+        &req.input_ids,
+        &req.target_mask,
+        &req.attention_mask,
+        batch_size,
+        seq_len,
+    )
+    .map_err(|e| err_resp(&e))?;
+    submit_ep_multi_lora(
+        state,
+        id,
+        headers,
+        tensors,
+        payload,
+        batch_size,
+        seq_len,
+        req.n_total,
+        req.lora_rank,
+        req.adapter_ids,
+        req.expected_steps,
+        req.allow_aggregate_loss,
+    )
+    .await
+}
 
-    let cmd = rustrain_ipc::EpCommand::TrainMultiLora {
+async fn ep_train_multi_lora_binary(
+    State(state): State<Arc<EpRouterState>>,
+    Path(id): Path<String>,
+    headers: HeaderMap,
+    body: Bytes,
+) -> Result<Json<TrainMultiLoraResponse>, (StatusCode, Json<ErrorResponse>)> {
+    let request = parse_binary_multi_lora_request(&body).map_err(|e| err_resp(&e))?;
+    crate::ep::validate_multi_lora_global_batch_size(
+        request.batch_size,
+        request.n_total,
+        state.world_size,
+    )
+    .map_err(|error| err_resp(&error))?;
+    submit_ep_multi_lora(
+        state,
+        id,
+        headers,
+        request.tensors,
+        request.payload,
+        request.batch_size,
+        request.seq_len,
+        request.n_total,
+        request.lora_rank,
+        request.adapter_ids,
+        request.expected_steps,
+        false,
+    )
+    .await
+}
+
+async fn submit_ep_multi_lora(
+    state: Arc<EpRouterState>,
+    id: String,
+    headers: HeaderMap,
+    tensors: rustrain_ipc::TensorSlabRef,
+    payload: Vec<u8>,
+    batch_size: usize,
+    seq_len: usize,
+    n_total: i32,
+    lora_rank: i32,
+    adapter_ids: Vec<i64>,
+    expected_steps: Vec<u64>,
+    allow_aggregate_loss: bool,
+) -> Result<Json<TrainMultiLoraResponse>, (StatusCode, Json<ErrorResponse>)> {
+    if n_total <= 0 || lora_rank <= 0 {
+        return Err(err_resp(
+            "multi-LoRA n_total and lora_rank must be positive",
+        ));
+    }
+    if !adapter_ids.is_empty() {
+        if adapter_ids.len() != n_total as usize {
+            return Err(err_resp(&format!(
+                "adapter_ids length {} must match n_total={}",
+                adapter_ids.len(),
+                n_total
+            )));
+        }
+        if adapter_ids.iter().any(|id| *id <= 0) {
+            return Err(err_resp("adapter_ids must contain only positive IDs"));
+        }
+        if adapter_ids.iter().copied().collect::<HashSet<_>>().len() != adapter_ids.len() {
+            return Err(err_resp("adapter_ids must not contain duplicates"));
+        }
+        if !expected_steps.is_empty() && expected_steps.len() != adapter_ids.len() {
+            return Err(err_resp(&format!(
+                "expected_steps length {} must match adapter_ids length {}",
+                expected_steps.len(),
+                adapter_ids.len()
+            )));
+        }
+    } else if !expected_steps.is_empty() {
+        return Err(err_resp("expected_steps requires adapter_ids"));
+    }
+    let source_count =
+        multi_lora_source_count(batch_size, n_total as usize).map_err(|error| err_resp(&error))?;
+    let normalized_payload_bytes =
+        normalized_multi_lora_payload_bytes(seq_len, n_total as usize, source_count)
+            .map_err(|error| err_resp(&error))?;
+    if !state.coordinator.is_healthy() {
+        return Err(ep_dispatch_unavailable("EP coordinator is unavailable"));
+    }
+    let (response, receiver) = oneshot::channel();
+    let request = MultiLoraDispatchRequest {
         session_id: id,
-        input_ids,
-        target_mask,
-        attention_mask,
-        seq_len,
-        n_total: req.n_total,
-        lora_rank: req.lora_rank,
+        tensors,
+        payload,
+        n_total: n_total as usize,
+        lora_rank,
+        adapter_ids,
+        expected_steps,
+        source_count,
+        normalized_payload_bytes,
+        // Capability v1 means the client understands per-adapter losses,
+        // optimizer steps, and the explicit coalesced loss scope returned
+        // below. Keep the body flag as a backwards-compatible override.
+        allow_aggregate_loss: allow_aggregate_loss || multi_lora_capability_v1(&headers),
+        response,
     };
-    match state.coordinator.dispatch(&cmd) {
-        rustrain_ipc::EpResult::Loss(loss) => Ok(Json(TrainStepResponse { loss, step: 0 })),
+    {
+        let _submission = state
+            .dispatch_submission
+            .lock()
+            .map_err(|_| ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed))?;
+        state
+            .multi_lora_batcher
+            .submit(&state.dispatcher, Arc::clone(&state.coordinator), request)
+            .map_err(ep_dispatch_schedule_error)?;
+    }
+    let outcome = receiver
+        .await
+        .map_err(|_| ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed))?;
+    match outcome.result {
+        rustrain_ipc::EpResult::MultiLoraTrain {
+            loss,
+            step,
+            adapter_losses,
+            adapter_steps,
+        } => Ok(Json(TrainMultiLoraResponse {
+            capability_version: 1,
+            capabilities: MULTI_LORA_RESPONSE_CAPABILITIES,
+            loss,
+            step,
+            loss_scope: if outcome.request_count > 1 {
+                "coalesced_batch"
+            } else {
+                "request"
+            },
+            coalesced_requests: outcome.request_count,
+            adapter_losses,
+            adapter_steps,
+        })),
+        rustrain_ipc::EpResult::Train { loss, step } => Ok(Json(TrainMultiLoraResponse {
+            capability_version: 1,
+            capabilities: MULTI_LORA_RESPONSE_CAPABILITIES,
+            loss,
+            step,
+            loss_scope: if outcome.request_count > 1 {
+                "coalesced_batch"
+            } else {
+                "request"
+            },
+            coalesced_requests: outcome.request_count,
+            adapter_losses: Vec::new(),
+            adapter_steps: Vec::new(),
+        })),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
     }
 }
 
+async fn dispatch_slab(
+    state: &EpRouterState,
+    command: rustrain_ipc::EpCommand,
+    payload: Vec<u8>,
+) -> Result<rustrain_ipc::EpResult, (StatusCode, Json<ErrorResponse>)> {
+    if !state.coordinator.is_healthy() {
+        return Err(ep_dispatch_unavailable("EP coordinator is unavailable"));
+    }
+    let coordinator = Arc::clone(&state.coordinator);
+    let receiver = {
+        let _submission = state
+            .dispatch_submission
+            .lock()
+            .map_err(|_| ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed))?;
+        state.multi_lora_batcher.seal_current();
+        state
+            .dispatcher
+            .submit(move || coordinator.dispatch_with_slab(&command, &payload))
+            .map_err(ep_dispatch_schedule_error)?
+    };
+    receiver
+        .await
+        .map_err(|_| ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed))
+}
+
+async fn dispatch_ep(
+    state: &EpRouterState,
+    command: rustrain_ipc::EpCommand,
+) -> Result<rustrain_ipc::EpResult, (StatusCode, Json<ErrorResponse>)> {
+    if !state.coordinator.is_healthy() {
+        return Err(ep_dispatch_unavailable("EP coordinator is unavailable"));
+    }
+    let coordinator = Arc::clone(&state.coordinator);
+    let receiver = {
+        let _submission = state
+            .dispatch_submission
+            .lock()
+            .map_err(|_| ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed))?;
+        state.multi_lora_batcher.seal_current();
+        state
+            .dispatcher
+            .submit(move || coordinator.dispatch(&command))
+            .map_err(ep_dispatch_schedule_error)?
+    };
+    receiver
+        .await
+        .map_err(|_| ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed))
+}
+
+fn ep_dispatch_schedule_error(error: EpDispatchScheduleError) -> (StatusCode, Json<ErrorResponse>) {
+    let status = match error {
+        EpDispatchScheduleError::QueueFull => StatusCode::TOO_MANY_REQUESTS,
+        EpDispatchScheduleError::QueueClosed | EpDispatchScheduleError::WorkerFailed => {
+            StatusCode::SERVICE_UNAVAILABLE
+        }
+    };
+    (
+        status,
+        Json(ErrorResponse {
+            error: error.to_string(),
+        }),
+    )
+}
+
+fn ep_dispatch_unavailable(message: &str) -> (StatusCode, Json<ErrorResponse>) {
+    (
+        StatusCode::SERVICE_UNAVAILABLE,
+        Json(ErrorResponse {
+            error: message.to_string(),
+        }),
+    )
+}
+
+#[cfg(test)]
+mod tensor_http_shape_tests {
+    use base64::{engine::general_purpose, Engine};
+
+    use super::{
+        build_multi_lora_window_payload, decode_int64_bytes, ep_dispatch_schedule_error,
+        multi_lora_capability_v1, multi_lora_source_count, normalized_multi_lora_payload_bytes,
+        pack_tensor_slab, parse_binary_multi_lora_request, project_multi_lora_result,
+        validate_multi_lora_http_shapes, validate_selected_eval_http_shapes,
+        validate_train_http_shapes, EpDispatchScheduleError, EpDispatchScheduler, HeaderMap,
+        HeaderValue, MultiLoraBatchConfig, MultiLoraBatchState, MultiLoraBatcher,
+        MultiLoraDispatchRequest, MultiLoraWindow, StatusCode, TensorHttp,
+        MULTI_LORA_CAPABILITY_HEADER,
+    };
+
+    fn tensor(shape: &[i64]) -> TensorHttp {
+        TensorHttp {
+            data: String::new(),
+            shape: shape.to_vec(),
+            dtype: "int64".into(),
+        }
+    }
+
+    fn encoded_tensor(rows: &[i64]) -> TensorHttp {
+        TensorHttp {
+            data: general_purpose::STANDARD.encode(
+                rows.iter()
+                    .copied()
+                    .flat_map(i64::to_le_bytes)
+                    .collect::<Vec<_>>(),
+            ),
+            shape: vec![rows.len() as i64, 1],
+            dtype: "int64".into(),
+        }
+    }
+
+    fn batch_request(
+        session_id: &str,
+        rows: &[i64],
+        n_total: usize,
+        adapter_ids: &[i64],
+    ) -> MultiLoraDispatchRequest {
+        let tensor = encoded_tensor(rows);
+        let (tensors, payload) =
+            pack_tensor_slab(&tensor, &tensor, &tensor, rows.len(), 1).unwrap();
+        let source_count = multi_lora_source_count(rows.len(), n_total).unwrap();
+        let normalized_payload_bytes =
+            normalized_multi_lora_payload_bytes(1, n_total, source_count).unwrap();
+        let (response, _receiver) = tokio::sync::oneshot::channel();
+        MultiLoraDispatchRequest {
+            session_id: session_id.to_string(),
+            tensors,
+            payload,
+            n_total,
+            lora_rank: 8,
+            adapter_ids: adapter_ids.to_vec(),
+            expected_steps: Vec::new(),
+            source_count,
+            normalized_payload_bytes,
+            allow_aggregate_loss: true,
+            response,
+        }
+    }
+
+    fn batch_request_with_rank(
+        session_id: &str,
+        rows: &[i64],
+        n_total: usize,
+        adapter_ids: &[i64],
+        lora_rank: i32,
+    ) -> MultiLoraDispatchRequest {
+        let mut request = batch_request(session_id, rows, n_total, adapter_ids);
+        request.lora_rank = lora_rank;
+        request
+    }
+
+    fn slab_i64(payload: &[u8], span: rustrain_ipc::TensorSpan) -> Vec<i64> {
+        let start = span.offset_bytes as usize;
+        let end = start + span.len_bytes as usize;
+        payload[start..end]
+            .chunks_exact(8)
+            .map(|bytes| i64::from_le_bytes(bytes.try_into().unwrap()))
+            .collect()
+    }
+
+    #[test]
+    fn train_shape_keeps_batch_and_sequence_dimensions() {
+        let input = tensor(&[8, 128]);
+        let target = tensor(&[8, 128]);
+        let attention = tensor(&[8, 128]);
+
+        assert_eq!(
+            validate_train_http_shapes(&input, &target, &attention).unwrap(),
+            (8, 128)
+        );
+    }
+
+    #[test]
+    fn multi_lora_capability_v1_accepts_versioned_header_tokens() {
+        let mut headers = HeaderMap::new();
+        assert!(!multi_lora_capability_v1(&headers));
+        headers.insert(MULTI_LORA_CAPABILITY_HEADER, HeaderValue::from_static("v1"));
+        assert!(multi_lora_capability_v1(&headers));
+        headers.insert(
+            MULTI_LORA_CAPABILITY_HEADER,
+            HeaderValue::from_static("v0, V1"),
+        );
+        assert!(multi_lora_capability_v1(&headers));
+        headers.insert(
+            MULTI_LORA_CAPABILITY_HEADER,
+            HeaderValue::from_static("v10"),
+        );
+        assert!(!multi_lora_capability_v1(&headers));
+    }
+
+    #[test]
+    fn binary_multi_lora_request_parses_wire_sections_without_base64() {
+        let mut bytes = Vec::new();
+        bytes.extend_from_slice(b"RLM1");
+        bytes.extend_from_slice(&1u16.to_le_bytes());
+        bytes.extend_from_slice(&0u16.to_le_bytes());
+        bytes.extend_from_slice(&2u32.to_le_bytes()); // batch
+        bytes.extend_from_slice(&2u32.to_le_bytes()); // sequence
+        bytes.extend_from_slice(&1u32.to_le_bytes()); // n_total
+        bytes.extend_from_slice(&4i32.to_le_bytes()); // LoRA rank
+        bytes.extend_from_slice(&1u32.to_le_bytes()); // adapter count
+        bytes.extend_from_slice(&1u32.to_le_bytes()); // expected-step count
+        for _ in 0..3 {
+            bytes.extend_from_slice(&32u64.to_le_bytes());
+        }
+        bytes.extend_from_slice(&7i64.to_le_bytes());
+        bytes.extend_from_slice(&3i64.to_le_bytes());
+        for value in 0i64..12 {
+            bytes.extend_from_slice(&value.to_le_bytes());
+        }
+        let request = parse_binary_multi_lora_request(&bytes).unwrap();
+        assert_eq!(request.batch_size, 2);
+        assert_eq!(request.seq_len, 2);
+        assert_eq!(request.adapter_ids, vec![7]);
+        assert_eq!(request.expected_steps, vec![3]);
+        assert_eq!(request.payload.len(), 160);
+        assert_eq!(request.tensors.input_ids.len_bytes, 32);
+        assert_eq!(request.tensors.target_mask.len_bytes, 32);
+        assert_eq!(request.tensors.attention_mask.len_bytes, 32);
+        request.tensors.validate(request.payload.len()).unwrap();
+    }
+
+    #[test]
+    fn binary_multi_lora_request_rejects_trailing_or_invalid_sections() {
+        let mut bytes = b"RLM1".to_vec();
+        bytes.extend_from_slice(&1u16.to_le_bytes());
+        bytes.extend_from_slice(&0u16.to_le_bytes());
+        bytes.extend_from_slice(&[0; 48]);
+        assert!(parse_binary_multi_lora_request(&bytes).is_err());
+        bytes[0] = b'X';
+        assert!(parse_binary_multi_lora_request(&bytes).is_err());
+    }
+
+    #[test]
+    fn train_shape_rejects_mismatched_masks_and_non_positive_dimensions() {
+        let input = tensor(&[4, 16]);
+        assert!(validate_train_http_shapes(&input, &tensor(&[2, 16]), &tensor(&[4, 16])).is_err());
+        assert!(validate_train_http_shapes(
+            &tensor(&[0, 16]),
+            &tensor(&[0, 16]),
+            &tensor(&[0, 16])
+        )
+        .is_err());
+        assert!(
+            validate_train_http_shapes(&tensor(&[4, 0]), &tensor(&[4, 0]), &tensor(&[4, 0]))
+                .is_err()
+        );
+    }
+
+    #[test]
+    fn multi_lora_shape_accepts_replica_and_global_row_layouts() {
+        for rows in [1, 3, 6] {
+            let input = tensor(&[rows, 32]);
+            assert_eq!(
+                validate_multi_lora_http_shapes(
+                    &input,
+                    &tensor(&[rows, 32]),
+                    &tensor(&[rows, 32]),
+                    3
+                )
+                .unwrap(),
+                (rows as usize, 32)
+            );
+        }
+        let invalid = tensor(&[5, 32]);
+        assert!(
+            validate_multi_lora_http_shapes(&invalid, &tensor(&[5, 32]), &tensor(&[5, 32]), 3)
+                .is_err()
+        );
+    }
+
+    #[test]
+    fn selected_eval_shape_requires_one_row_or_one_row_per_adapter() {
+        let one = tensor(&[1, 32]);
+        assert_eq!(
+            validate_selected_eval_http_shapes(&one, &one, &one, &[11, 12]).unwrap(),
+            (1, 32)
+        );
+        let two = tensor(&[2, 32]);
+        assert_eq!(
+            validate_selected_eval_http_shapes(&two, &two, &two, &[11, 12]).unwrap(),
+            (2, 32)
+        );
+        assert!(validate_selected_eval_http_shapes(&two, &two, &two, &[11]).is_err());
+        assert!(validate_selected_eval_http_shapes(&one, &one, &one, &[0, 12]).is_err());
+    }
+
+    #[test]
+    fn coalesced_multi_lora_preserves_source_major_adapter_rows() {
+        let mut window = MultiLoraWindow::new(batch_request(
+            "tenant-session",
+            &[10, 11, 20, 21],
+            2,
+            &[101, 102],
+        ));
+        window.push(batch_request("tenant-session", &[12, 22], 1, &[103]));
+
+        let (tensors, payload, adapter_ids, _, _) =
+            build_multi_lora_window_payload(&window).unwrap();
+        assert_eq!(adapter_ids, vec![101, 102, 103]);
+        assert_eq!(tensors.batch_size, 6);
+        assert_eq!(
+            slab_i64(&payload, tensors.input_ids),
+            vec![10, 11, 12, 20, 21, 22]
+        );
+        assert_eq!(
+            slab_i64(&payload, tensors.target_mask),
+            vec![10, 11, 12, 20, 21, 22]
+        );
+        tensors.validate(payload.len()).unwrap();
+    }
+
+    #[test]
+    fn coalesced_multi_lora_expands_request_shared_rows_per_adapter() {
+        let mut window =
+            MultiLoraWindow::new(batch_request("tenant-session", &[7], 2, &[101, 102]));
+        window.push(batch_request("tenant-session", &[8], 1, &[103]));
+
+        let (tensors, payload, adapter_ids, _, _) =
+            build_multi_lora_window_payload(&window).unwrap();
+        assert_eq!(adapter_ids, vec![101, 102, 103]);
+        assert_eq!(slab_i64(&payload, tensors.input_ids), vec![7, 7, 8]);
+    }
+
+    #[test]
+    fn coalesced_multi_lora_preserves_expected_optimizer_steps() {
+        let mut first = batch_request("tenant-session", &[7], 2, &[101, 102]);
+        first.expected_steps = vec![3, 5];
+        let mut window = MultiLoraWindow::new(first);
+        let mut second = batch_request("tenant-session", &[8], 1, &[103]);
+        second.expected_steps = vec![9];
+        window.push(second);
+
+        let (_, _, adapter_ids, _, expected_steps) =
+            build_multi_lora_window_payload(&window).unwrap();
+        assert_eq!(adapter_ids, vec![101, 102, 103]);
+        assert_eq!(expected_steps, vec![3, 5, 9]);
+    }
+
+    #[test]
+    fn coalesced_result_projects_only_the_request_adapter_range() {
+        let result = rustrain_ipc::EpResult::MultiLoraTrain {
+            loss: 2.0,
+            step: 9,
+            adapter_losses: vec![
+                rustrain_ipc::command::AdapterLoss {
+                    adapter_id: 101,
+                    loss: 1.0,
+                },
+                rustrain_ipc::command::AdapterLoss {
+                    adapter_id: 102,
+                    loss: 2.0,
+                },
+                rustrain_ipc::command::AdapterLoss {
+                    adapter_id: 103,
+                    loss: 3.0,
+                },
+            ],
+            adapter_steps: vec![
+                rustrain_ipc::command::AdapterStep {
+                    adapter_id: 101,
+                    step: 4,
+                },
+                rustrain_ipc::command::AdapterStep {
+                    adapter_id: 102,
+                    step: 5,
+                },
+                rustrain_ipc::command::AdapterStep {
+                    adapter_id: 103,
+                    step: 6,
+                },
+            ],
+        };
+        let projected = project_multi_lora_result(&result, 1..3);
+        match projected {
+            rustrain_ipc::EpResult::MultiLoraTrain {
+                loss,
+                step,
+                adapter_losses,
+                adapter_steps,
+            } => {
+                assert_eq!(loss, 2.0);
+                assert_eq!(step, 9);
+                assert_eq!(
+                    adapter_losses
+                        .iter()
+                        .map(|item| item.adapter_id)
+                        .collect::<Vec<_>>(),
+                    vec![102, 103]
+                );
+                assert_eq!(
+                    adapter_steps
+                        .iter()
+                        .map(|item| (item.adapter_id, item.step))
+                        .collect::<Vec<_>>(),
+                    vec![(102, 5), (103, 6)]
+                );
+            }
+            other => panic!("unexpected projected result: {other:?}"),
+        }
+
+        let legacy = rustrain_ipc::EpResult::MultiLoraTrain {
+            loss: 4.0,
+            step: 10,
+            adapter_losses: Vec::new(),
+            adapter_steps: Vec::new(),
+        };
+        assert!(matches!(
+            project_multi_lora_result(&legacy, 0..1),
+            rustrain_ipc::EpResult::Train {
+                loss: 4.0,
+                step: 10
+            }
+        ));
+    }
+
+    #[test]
+    fn coalescing_rejects_overlapping_tenants_and_capacity_overflow() {
+        let config = MultiLoraBatchConfig {
+            window: std::time::Duration::from_millis(1),
+            max_requests: 2,
+            max_open_windows: 2,
+            max_adapters: 3,
+            max_rank_work: 64,
+            max_payload_bytes: usize::MAX,
+        };
+        let mut window = MultiLoraWindow::new(batch_request("session", &[1], 1, &[11]));
+        assert!(!window.can_accept(&batch_request("session", &[2], 1, &[11]), config));
+        assert!(!window.can_accept(&batch_request("other", &[2], 1, &[12]), config));
+        assert!(window.can_accept(
+            &batch_request_with_rank("session", &[2], 1, &[12], 16),
+            config
+        ));
+        let tight_config = MultiLoraBatchConfig {
+            max_rank_work: 16,
+            ..config
+        };
+        assert!(!window.can_accept(
+            &batch_request_with_rank("session", &[2], 1, &[12], 16),
+            tight_config
+        ));
+        let mut guarded = batch_request("session", &[2], 1, &[12]);
+        guarded.expected_steps = vec![0];
+        assert!(!window.can_accept(&guarded, config));
+        assert!(window.can_accept(&batch_request("session", &[2, 3], 2, &[12, 13]), config));
+        window.push(batch_request("session", &[2], 1, &[12]));
+        assert!(!window.can_accept(&batch_request("session", &[3], 1, &[13]), config));
+    }
+
+    #[test]
+    fn expected_steps_make_request_exclusive() {
+        let config = MultiLoraBatchConfig {
+            window: std::time::Duration::from_millis(1),
+            max_requests: 4,
+            max_open_windows: 2,
+            max_adapters: 4,
+            max_rank_work: 64,
+            max_payload_bytes: usize::MAX,
+        };
+        let mut first = batch_request("session", &[1], 1, &[11]);
+        first.expected_steps = vec![0];
+        let second = batch_request("session", &[2], 1, &[12]);
+        assert!(!first.coalescible());
+        assert!(!MultiLoraWindow::new(first).can_accept(&second, config));
+
+        let mut second_guarded = batch_request("session", &[2], 1, &[12]);
+        second_guarded.expected_steps = vec![0];
+        assert!(!second_guarded.coalescible());
+    }
+
+    #[test]
+    fn coalescing_wait_wakes_when_open_window_is_sealed() {
+        let config = MultiLoraBatchConfig {
+            window: std::time::Duration::from_secs(2),
+            max_requests: 4,
+            max_open_windows: 2,
+            max_adapters: 4,
+            max_rank_work: 64,
+            max_payload_bytes: usize::MAX,
+        };
+        let batcher = MultiLoraBatcher::new(config);
+        let window_id = 7;
+        batcher
+            .state
+            .lock()
+            .unwrap()
+            .open_window_ids
+            .push_back(window_id);
+
+        let waiting_batcher = batcher.clone();
+        let (started, started_rx) = std::sync::mpsc::channel();
+        let waiter = std::thread::spawn(move || {
+            let start = std::time::Instant::now();
+            started.send(()).unwrap();
+            waiting_batcher.wait_for_window(
+                window_id,
+                std::time::Instant::now() + config.window,
+            );
+            start.elapsed()
+        });
+        started_rx.recv().unwrap();
+        std::thread::sleep(std::time::Duration::from_millis(10));
+        batcher.seal_current();
+
+        assert!(
+            waiter.join().unwrap() < std::time::Duration::from_millis(500),
+            "sealed coalescing window must not wait for its full deadline"
+        );
+    }
+
+    #[test]
+    fn coalescing_wait_wakes_when_window_reaches_capacity() {
+        let config = MultiLoraBatchConfig {
+            window: std::time::Duration::from_secs(2),
+            max_requests: 1,
+            max_open_windows: 2,
+            max_adapters: 4,
+            max_rank_work: 64,
+            max_payload_bytes: usize::MAX,
+        };
+        let batcher = MultiLoraBatcher::new(config);
+        let window_id = 11;
+        {
+            let mut state = batcher.state.lock().unwrap();
+            state.open_window_ids.push_back(window_id);
+            state.windows.insert(
+                window_id,
+                MultiLoraWindow::new(batch_request("session", &[1], 1, &[11])),
+            );
+        }
+
+        let waiting_batcher = batcher.clone();
+        let (started, started_rx) = std::sync::mpsc::channel();
+        let waiter = std::thread::spawn(move || {
+            let start = std::time::Instant::now();
+            started.send(()).unwrap();
+            waiting_batcher.wait_for_window(
+                window_id,
+                std::time::Instant::now() + config.window,
+            );
+            start.elapsed()
+        });
+        started_rx.recv().unwrap();
+        std::thread::sleep(std::time::Duration::from_millis(10));
+        {
+            let mut state = batcher.state.lock().unwrap();
+            batcher.seal_window_if_at_capacity(&mut state, window_id);
+        }
+
+        assert!(
+            waiter.join().unwrap() < std::time::Duration::from_millis(500),
+            "full coalescing window must not wait for its full deadline"
+        );
+        assert!(batcher.state.lock().unwrap().open_window_ids.is_empty());
+    }
+
+    #[test]
+    fn coalescing_keeps_interleaved_layout_buckets_open() {
+        let config = MultiLoraBatchConfig {
+            window: std::time::Duration::from_millis(30),
+            max_requests: 4,
+            max_open_windows: 2,
+            max_adapters: 4,
+            max_rank_work: 64,
+            max_payload_bytes: usize::MAX,
+        };
+        let mut state = MultiLoraBatchState::default();
+        state.open_window_ids.extend([17, 19]);
+        state.windows.insert(
+            17,
+            MultiLoraWindow::new(batch_request("session", &[1], 1, &[11])),
+        );
+        state.windows.insert(
+            19,
+            MultiLoraWindow::new(batch_request("session", &[2, 3], 1, &[21])),
+        );
+
+        assert_eq!(
+            MultiLoraBatcher::compatible_window_id(
+                &state,
+                &batch_request("session", &[4], 1, &[12]),
+                config,
+            ),
+            Some(17)
+        );
+        assert_eq!(
+            MultiLoraBatcher::compatible_window_id(
+                &state,
+                &batch_request("session", &[5, 6], 1, &[22]),
+                config,
+            ),
+            Some(19)
+        );
+        assert_eq!(
+            state.open_window_ids.iter().copied().collect::<Vec<_>>(),
+            vec![17, 19]
+        );
+    }
+
+    #[test]
+    fn coalescing_wait_keeps_the_deadline_for_an_open_window() {
+        let config = MultiLoraBatchConfig {
+            window: std::time::Duration::from_millis(30),
+            max_requests: 4,
+            max_open_windows: 2,
+            max_adapters: 4,
+            max_rank_work: 64,
+            max_payload_bytes: usize::MAX,
+        };
+        let batcher = MultiLoraBatcher::new(config);
+        let window_id = 13;
+        batcher
+            .state
+            .lock()
+            .unwrap()
+            .open_window_ids
+            .push_back(window_id);
+
+        let start = std::time::Instant::now();
+        batcher.wait_for_window(window_id, std::time::Instant::now() + config.window);
+        assert!(
+            start.elapsed() >= std::time::Duration::from_millis(20),
+            "open coalescing window must retain its normal batching deadline"
+        );
+        assert!(batcher.state.lock().unwrap().open_window_ids.is_empty());
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn coalescing_deadline_overlaps_a_busy_dispatcher() {
+        let config = MultiLoraBatchConfig {
+            window: std::time::Duration::from_millis(40),
+            max_requests: 4,
+            max_open_windows: 2,
+            max_adapters: 4,
+            max_rank_work: 64,
+            max_payload_bytes: usize::MAX,
+        };
+        let batcher = MultiLoraBatcher::new(config);
+        let window_id = 23;
+        batcher
+            .state
+            .lock()
+            .unwrap()
+            .open_window_ids
+            .push_back(window_id);
+
+        let scheduler = EpDispatchScheduler::new(2);
+        let gate = std::sync::Arc::new((
+            std::sync::Mutex::new(false),
+            std::sync::Condvar::new(),
+        ));
+        let job_gate = std::sync::Arc::clone(&gate);
+        let (started, started_rx) = tokio::sync::oneshot::channel();
+        let first = scheduler
+            .submit(move || {
+                let _ = started.send(());
+                let (lock, wake) = &*job_gate;
+                let mut released = lock.lock().unwrap();
+                while !*released {
+                    released = wake.wait(released).unwrap();
+                }
+            })
+            .unwrap();
+        started_rx.await.unwrap();
+
+        let deadline = std::time::Instant::now() + config.window;
+        let waiting_batcher = batcher.clone();
+        let second = scheduler
+            .submit(move || {
+                let start = std::time::Instant::now();
+                waiting_batcher.wait_for_window(window_id, deadline);
+                start.elapsed()
+            })
+            .unwrap();
+
+        tokio::time::sleep(std::time::Duration::from_millis(80)).await;
+        {
+            let (lock, wake) = &*gate;
+            *lock.lock().unwrap() = true;
+            wake.notify_all();
+        }
+        first.await.unwrap();
+        assert!(
+            second.await.unwrap() < std::time::Duration::from_millis(20),
+            "an expired batching deadline must dispatch immediately after the previous GPU job"
+        );
+        assert!(batcher.state.lock().unwrap().open_window_ids.is_empty());
+    }
+
+    #[test]
+    fn implicit_registry_request_is_an_exclusive_window() {
+        let config = MultiLoraBatchConfig {
+            window: std::time::Duration::from_millis(1),
+            max_requests: 4,
+            max_open_windows: 2,
+            max_adapters: 4,
+            max_rank_work: 64,
+            max_payload_bytes: usize::MAX,
+        };
+        let window = MultiLoraWindow::new(batch_request("session", &[1], 1, &[]));
+        assert!(!window.can_accept(&batch_request("session", &[2], 1, &[12]), config));
+    }
+
+    #[test]
+    fn request_local_loss_opt_out_is_an_exclusive_window() {
+        let config = MultiLoraBatchConfig {
+            window: std::time::Duration::from_millis(1),
+            max_requests: 4,
+            max_open_windows: 2,
+            max_adapters: 4,
+            max_rank_work: 64,
+            max_payload_bytes: usize::MAX,
+        };
+        let mut request = batch_request("session", &[1], 1, &[11]);
+        request.allow_aggregate_loss = false;
+        let window = MultiLoraWindow::new(request);
+        assert!(!window.can_accept(&batch_request("session", &[2], 1, &[12]), config));
+    }
+
+    #[test]
+    fn int64_decode_rejects_wrong_dtype_and_trailing_bytes() {
+        let mut value = tensor(&[1]);
+        value.dtype = "float32".into();
+        value.data = general_purpose::STANDARD.encode(1_i64.to_le_bytes());
+        assert!(decode_int64_bytes(&value).is_err());
+
+        value.dtype = "int64".into();
+        value.data = general_purpose::STANDARD.encode([0_u8; 9]);
+        assert!(decode_int64_bytes(&value).is_err());
+    }
+
+    #[test]
+    fn tensor_slab_pack_aligns_spans_and_preserves_wire_bytes() {
+        let wire = [1_i64, -2_i64]
+            .into_iter()
+            .flat_map(i64::to_le_bytes)
+            .collect::<Vec<_>>();
+        let encoded = general_purpose::STANDARD.encode(&wire);
+        let make = || TensorHttp {
+            data: encoded.clone(),
+            shape: vec![1, 2],
+            dtype: "int64".into(),
+        };
+        let (reference, payload) = pack_tensor_slab(&make(), &make(), &make(), 1, 2).unwrap();
+        for span in reference.spans() {
+            assert_eq!(span.offset_bytes % 64, 0);
+            let start = span.offset_bytes as usize;
+            let end = start + span.len_bytes as usize;
+            assert_eq!(&payload[start..end], wire);
+        }
+        reference.validate(payload.len()).unwrap();
+    }
+
+    #[test]
+    fn dispatch_pressure_and_worker_failure_have_distinct_statuses() {
+        assert_eq!(
+            ep_dispatch_schedule_error(EpDispatchScheduleError::QueueFull).0,
+            StatusCode::TOO_MANY_REQUESTS
+        );
+        assert_eq!(
+            ep_dispatch_schedule_error(EpDispatchScheduleError::QueueClosed).0,
+            StatusCode::SERVICE_UNAVAILABLE
+        );
+        assert_eq!(
+            ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed).0,
+            StatusCode::SERVICE_UNAVAILABLE
+        );
+    }
+}
+
 async fn ep_add_lora(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
     Json(req): Json<AddLoRAHttp>,
 ) -> Result<Json<AddLoRAResponse>, (StatusCode, Json<ErrorResponse>)> {
@@ -737,8 +2997,12 @@ async fn ep_add_lora(
         alpha: req.alpha,
         target_layers: req.target_layers,
         target_modules: req.target_modules,
+        optimizer_lr: req.optimizer_lr,
+        optimizer_beta1: req.optimizer_beta1,
+        optimizer_beta2: req.optimizer_beta2,
+        optimizer_eps: req.optimizer_eps,
     };
-    match state.coordinator.dispatch(&cmd) {
+    match dispatch_ep(&state, cmd).await? {
         rustrain_ipc::EpResult::AdapterId(id) => Ok(Json(AddLoRAResponse { adapter_id: id })),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
@@ -752,13 +3016,22 @@ struct BatchAddLoRAHttp {
     alpha: f64,
     target_layers: Vec<i64>,
     target_modules: String,
+    #[serde(default)]
+    optimizer_lr: Option<f64>,
+    #[serde(default)]
+    optimizer_beta1: Option<f64>,
+    #[serde(default)]
+    optimizer_beta2: Option<f64>,
+    #[serde(default)]
+    optimizer_eps: Option<f64>,
 }
 
 async fn ep_batch_add_lora(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
     Json(req): Json<BatchAddLoRAHttp>,
 ) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
+    crate::ep::validate_batch_add_lora_count(req.count).map_err(|error| err_resp(&error))?;
     let cmd = rustrain_ipc::EpCommand::BatchAddLora {
         session_id: id,
         count: req.count,
@@ -766,8 +3039,12 @@ async fn ep_batch_add_lora(
         alpha: req.alpha,
         target_layers: req.target_layers,
         target_modules: req.target_modules,
+        optimizer_lr: req.optimizer_lr,
+        optimizer_beta1: req.optimizer_beta1,
+        optimizer_beta2: req.optimizer_beta2,
+        optimizer_eps: req.optimizer_eps,
     };
-    match state.coordinator.dispatch(&cmd) {
+    match dispatch_ep(&state, cmd).await? {
         rustrain_ipc::EpResult::Count(n) => Ok(Json(serde_json::json!({"count": n}))),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
@@ -775,12 +3052,15 @@ async fn ep_batch_add_lora(
 }
 
 async fn ep_remove_lora(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
     Json(req): Json<RemoveLoRAHttp>,
 ) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
-    let cmd = rustrain_ipc::EpCommand::RemoveLora { session_id: id, adapter_id: req.adapter_id };
-    match state.coordinator.dispatch(&cmd) {
+    let cmd = rustrain_ipc::EpCommand::RemoveLora {
+        session_id: id,
+        adapter_id: req.adapter_id,
+    };
+    match dispatch_ep(&state, cmd).await? {
         rustrain_ipc::EpResult::Ok => Ok(Json(serde_json::json!({"removed": true}))),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
@@ -788,11 +3068,11 @@ async fn ep_remove_lora(
 }
 
 async fn ep_list_lora(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
 ) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
     let cmd = rustrain_ipc::EpCommand::ListLora { session_id: id };
-    match state.coordinator.dispatch(&cmd) {
+    match dispatch_ep(&state, cmd).await? {
         rustrain_ipc::EpResult::AdapterIds(ids) => Ok(Json(serde_json::json!({"adapters": ids}))),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
@@ -800,27 +3080,110 @@ async fn ep_list_lora(
 }
 
 async fn ep_export_adapter(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
     Json(req): Json<ExportHttp>,
 ) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
-    let cmd = rustrain_ipc::EpCommand::ExportAdapter { session_id: id, path: req.path };
-    match state.coordinator.dispatch(&cmd) {
+    let generation = state
+        .coordinator
+        .next_generation("export")
+        .map_err(|error| err_resp(&error))?;
+    let cmd = rustrain_ipc::EpCommand::ExportAdapter {
+        session_id: id,
+        path: req.path,
+        adapter_id: req.adapter_id,
+        generation,
+    };
+    match dispatch_ep(&state, cmd).await? {
         rustrain_ipc::EpResult::Count(n) => Ok(Json(serde_json::json!({"exported": n}))),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
     }
 }
 
+async fn ep_save_checkpoint(
+    State(state): State<Arc<EpRouterState>>,
+    Path(id): Path<String>,
+    Json(req): Json<CheckpointHttp>,
+) -> Result<Json<CheckpointResponse>, (StatusCode, Json<ErrorResponse>)> {
+    let path = req.path;
+    let response_path = path.clone();
+    if !state.coordinator.is_healthy() {
+        return Err(ep_dispatch_unavailable("EP coordinator is unavailable"));
+    }
+    let coordinator = Arc::clone(&state.coordinator);
+    let receiver = {
+        let _submission = state
+            .dispatch_submission
+            .lock()
+            .map_err(|_| ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed))?;
+        state.multi_lora_batcher.seal_current();
+        state
+            .dispatcher
+            .submit(move || coordinator.coordinated_save_checkpoint(&id, &path))
+            .map_err(ep_dispatch_schedule_error)?
+    };
+    let (step, loss) = receiver
+        .await
+        .map_err(|_| ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed))?
+        .map_err(|error| err_resp(&error))?;
+    Ok(Json(CheckpointResponse {
+        step,
+        loss,
+        path: response_path,
+    }))
+}
+
+async fn ep_load_checkpoint(
+    State(state): State<Arc<EpRouterState>>,
+    Path(id): Path<String>,
+    Json(req): Json<CheckpointHttp>,
+) -> Result<Json<CheckpointResponse>, (StatusCode, Json<ErrorResponse>)> {
+    let path = req.path;
+    let response_path = path.clone();
+    if !state.coordinator.is_healthy() {
+        return Err(ep_dispatch_unavailable("EP coordinator is unavailable"));
+    }
+    let coordinator = Arc::clone(&state.coordinator);
+    let receiver = {
+        let _submission = state
+            .dispatch_submission
+            .lock()
+            .map_err(|_| ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed))?;
+        state.multi_lora_batcher.seal_current();
+        state
+            .dispatcher
+            .submit(move || coordinator.coordinated_load_checkpoint(&id, &path))
+            .map_err(ep_dispatch_schedule_error)?
+    };
+    let (step, loss) = receiver
+        .await
+        .map_err(|_| ep_dispatch_schedule_error(EpDispatchScheduleError::WorkerFailed))?
+        .map_err(|error| err_resp(&error))?;
+    Ok(Json(CheckpointResponse {
+        step,
+        loss,
+        path: response_path,
+    }))
+}
+
 async fn ep_get_status(
-    State(state): State<Arc<EpAppState>>,
+    State(state): State<Arc<EpRouterState>>,
     Path(id): Path<String>,
 ) -> Result<Json<StatusResponse>, (StatusCode, Json<ErrorResponse>)> {
     let cmd = rustrain_ipc::EpCommand::Status { session_id: id };
-    match state.coordinator.dispatch(&cmd) {
-        rustrain_ipc::EpResult::Status { state, step, last_loss, model_path } => {
-            Ok(Json(StatusResponse { state, step, last_loss, model_path }))
-        }
+    match dispatch_ep(&state, cmd).await? {
+        rustrain_ipc::EpResult::Status {
+            state,
+            step,
+            last_loss,
+            model_path,
+        } => Ok(Json(StatusResponse {
+            state,
+            step,
+            last_loss,
+            model_path,
+        })),
         rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)),
         _ => Err(err_resp("unexpected result")),
     }
diff --git a/crates/rustrain-server/src/checkpoint.rs b/crates/rustrain-server/src/checkpoint.rs
index c3164fa1..82251d4f 100644
--- a/crates/rustrain-server/src/checkpoint.rs
+++ b/crates/rustrain-server/src/checkpoint.rs
@@ -1,203 +1,3 @@
-//! Checkpoint save/load: adapter (LoRA A/B) + optimizer state (Adam m/v) + step count.
+//! Compatibility facade for the shared Qwen checkpoint implementation.
 
-use anyhow::{Context, Result};
-use serde::{Deserialize, Serialize};
-use std::path::Path;
-use tch::Tensor;
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct CheckpointManifest {
-    pub format: String,
-    pub step: u64,
-    pub loss: f64,
-    pub model_path: String,
-    pub lora_rank: i64,
-    pub lora_alpha: i64,
-    pub files: Vec<String>,
-}
-
-pub struct CheckpointData {
-    pub manifest: CheckpointManifest,
-    pub lora_a: Vec<Tensor>,
-    pub lora_b: Vec<Tensor>,
-    pub adam_m: Vec<Tensor>,
-    pub adam_v: Vec<Tensor>,
-}
-
-/// Save checkpoint to a directory.
-/// Creates: manifest.json, adapter.safetensors, optimizer.safetensors
-pub fn save_checkpoint(
-    dir: &Path,
-    step: u64,
-    loss: f64,
-    model_path: &str,
-    lora_rank: i64,
-    lora_alpha: i64,
-    lora_a: &[Tensor],
-    lora_b: &[Tensor],
-    adam_m: &[Tensor],
-    adam_v: &[Tensor],
-) -> Result<()> {
-    std::fs::create_dir_all(dir)
-        .with_context(|| format!("create checkpoint dir {}", dir.display()))?;
-
-    // Save adapter (LoRA A/B) as safetensors
-    let adapter_path = dir.join("adapter.safetensors");
-    save_tensors(&adapter_path, &lora_a, &lora_b)?;
-
-    // Save optimizer state (Adam m/v) as safetensors
-    let optimizer_path = dir.join("optimizer.safetensors");
-    save_tensors(&optimizer_path, &adam_m, &adam_v)?;
-
-    // Write manifest
-    let manifest = CheckpointManifest {
-        format: "rustrain-checkpoint-v1".to_string(),
-        step,
-        loss,
-        model_path: model_path.to_string(),
-        lora_rank,
-        lora_alpha,
-        files: vec!["adapter.safetensors".into(), "optimizer.safetensors".into()],
-    };
-    let manifest_path = dir.join("manifest.json");
-    std::fs::write(&manifest_path, serde_json::to_string_pretty(&manifest)?)
-        .with_context(|| "write manifest.json")?;
-
-    tracing::info!(
-        step, loss,
-        path = dir.display().to_string(),
-        "checkpoint saved"
-    );
-    Ok(())
-}
-
-/// Load checkpoint from a directory.
-pub fn load_checkpoint(dir: &Path) -> Result<CheckpointData> {
-    let manifest_path = dir.join("manifest.json");
-    let manifest: CheckpointManifest = serde_json::from_str(
-        &std::fs::read_to_string(&manifest_path)
-            .with_context(|| format!("read {}", manifest_path.display()))?,
-    )
-    .with_context(|| "parse manifest.json")?;
-
-    let adapter_path = dir.join("adapter.safetensors");
-    let (lora_a, lora_b) = load_tensors(&adapter_path)?;
-
-    let optimizer_path = dir.join("optimizer.safetensors");
-    let (adam_m, adam_v) = load_tensors(&optimizer_path)?;
-
-    tracing::info!(
-        step = manifest.step,
-        loss = manifest.loss,
-        "checkpoint loaded"
-    );
-
-    Ok(CheckpointData {
-        manifest,
-        lora_a,
-        lora_b,
-        adam_m,
-        adam_v,
-    })
-}
-
-fn save_tensors(path: &Path, a: &[Tensor], b: &[Tensor]) -> Result<()> {
-    use std::io::Write;
-    let mut named: Vec<(String, Tensor)> = Vec::new();
-    for (i, t) in a.iter().enumerate() {
-        named.push((format!("a_{i}"), t.to_kind(tch::Kind::Float).to_device(tch::Device::Cpu)));
-    }
-    for (i, t) in b.iter().enumerate() {
-        named.push((format!("b_{i}"), t.to_kind(tch::Kind::Float).to_device(tch::Device::Cpu)));
-    }
-
-    // Build safetensors manually (header + data)
-    let mut header = serde_json::Map::new();
-    let mut offset = 0u64;
-    let mut all_bytes: Vec<u8> = Vec::new();
-    for (name, t) in &named {
-        let t = t.contiguous().to_kind(tch::Kind::Float);
-        let shape: Vec<i64> = t.size().iter().copied().collect();
-        let data: Vec<f32> = Vec::<f32>::try_from(&t.reshape([-1]))?;
-        let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
-        header.insert(
-            name.clone(),
-            serde_json::json!({"dtype":"F32","shape":shape,"data_offsets":[offset, offset + bytes.len() as u64]}),
-        );
-        offset += bytes.len() as u64;
-        all_bytes.extend_from_slice(&bytes);
-    }
-    let header_str = serde_json::to_string(&serde_json::Value::Object(header))?;
-    let file = std::fs::File::create(path).with_context(|| format!("create {}", path.display()))?;
-    let mut writer = std::io::BufWriter::new(file);
-    writer.write_all(&(header_str.len() as u64).to_le_bytes())?;
-    writer.write_all(header_str.as_bytes())?;
-    writer.write_all(&all_bytes)?;
-    Ok(())
-}
-
-fn load_tensors(path: &Path) -> Result<(Vec<Tensor>, Vec<Tensor>)> {
-    let data = std::fs::read(path).with_context(|| format!("read {}", path.display()))?;
-    if data.len() < 8 {
-        anyhow::bail!("safetensors file too small");
-    }
-    let header_len = u64::from_le_bytes(data[..8].try_into().unwrap()) as usize;
-    let header_str = std::str::from_utf8(&data[8..8 + header_len])?;
-    let header: serde_json::Value = serde_json::from_str(header_str)?;
-
-    // Count a_N and b_N entries
-    let mut a_count = 0;
-    let mut b_count = 0;
-    if let serde_json::Value::Object(map) = &header {
-        for key in map.keys() {
-            if key.starts_with("a_") {
-                a_count += 1;
-            } else if key.starts_with("b_") {
-                b_count += 1;
-            }
-        }
-    }
-
-    let mut a_tensors = Vec::new();
-    let mut b_tensors = Vec::new();
-
-    for i in 0..a_count {
-        let key = format!("a_{i}");
-        let (tensor, shape) = load_one_tensor(&header, &key, &data, 8 + header_len)?;
-        a_tensors.push(tensor.reshape(&shape));
-    }
-    for i in 0..b_count {
-        let key = format!("b_{i}");
-        let (tensor, shape) = load_one_tensor(&header, &key, &data, 8 + header_len)?;
-        b_tensors.push(tensor.reshape(&shape));
-    }
-
-    Ok((a_tensors, b_tensors))
-}
-
-fn load_one_tensor(
-    header: &serde_json::Value,
-    key: &str,
-    data: &[u8],
-    data_offset: usize,
-) -> Result<(Tensor, Vec<i64>)> {
-    let entry = header
-        .get(key)
-        .ok_or_else(|| anyhow::anyhow!("key {key} not found in safetensors"))?;
-    let shape: Vec<i64> = entry["shape"]
-        .as_array()
-        .unwrap()
-        .iter()
-        .map(|v| v.as_i64().unwrap())
-        .collect();
-    let offsets = entry["data_offsets"].as_array().unwrap();
-    let start = data_offset + offsets[0].as_u64().unwrap() as usize;
-    let end = data_offset + offsets[1].as_u64().unwrap() as usize;
-    let bytes = &data[start..end];
-    let floats: Vec<f32> = bytes
-        .chunks_exact(4)
-        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
-        .collect();
-    let tensor = Tensor::from_slice(&floats);
-    Ok((tensor, shape))
-}
+pub use rustrain_qwen3_6::checkpoint::*;
diff --git a/crates/rustrain-server/src/ep.rs b/crates/rustrain-server/src/ep.rs
index e84a5f35..258bf835 100644
--- a/crates/rustrain-server/src/ep.rs
+++ b/crates/rustrain-server/src/ep.rs
@@ -7,24 +7,66 @@
 //!
 //! Workers use `worker_main()` to enter the wait loop.
 
-use std::collections::HashMap;
+use std::ffi::CString;
 use std::io;
-use std::path::PathBuf;
-use std::sync::Arc;
+use std::ops::Range;
+use std::os::unix::ffi::OsStrExt;
+use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
 
-use rustrain_ipc::{EpChannel, EpCommand, EpResult, EpWorker};
+use rustrain_ipc::{EpChannel, EpCommand, EpResult, EpWorker, TensorSlabRef};
+use rustrain_parallel::topology::ParallelTopology;
 use tch::{Device, Kind};
 
 use crate::session::{
-    AddLoRARequest, EvalOutput, InitLoRARequest, Qwen36Session, SessLoadDatasetRequest,
-    SessLoadModelRequest, TrainInput, TrainOutput, TrainingSession,
+    configured_dynamic_pipeline_microbatches, AddLoRARequest, EvalOutput, InitLoRARequest,
+    MultiLoraEvalOutput, Qwen36Session, SessLoadDatasetRequest, SessLoadModelRequest, TrainOutput,
+    TrainingSession,
 };
 
+const TERMINAL_NATIVE_CONTEXT_PREFIX: &str = "terminal native LoRA context";
+
+fn is_terminal_native_context_result(result: &EpResult) -> bool {
+    matches!(
+        result,
+        EpResult::Error(error) if error.starts_with(TERMINAL_NATIVE_CONTEXT_PREFIX)
+    )
+}
+
+const DEFAULT_MAX_BATCH_ADD_LORA: usize = 64;
+const HARD_MAX_BATCH_ADD_LORA: usize = 4096;
+
+/// Validate before allocating the request-sized vector or touching CUDA.
+/// Keep the same bound in the HTTP process and every worker process so an
+/// invalid batch is rejected before the collective registry can change.
+pub fn validate_batch_add_lora_count(count: i32) -> Result<usize, String> {
+    if count <= 0 {
+        return Err("batch LoRA count must be positive".to_string());
+    }
+    let configured = std::env::var("RUSTRAIN_MAX_BATCH_ADD_LORA")
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .filter(|value| *value > 0)
+        .unwrap_or(DEFAULT_MAX_BATCH_ADD_LORA)
+        .min(HARD_MAX_BATCH_ADD_LORA);
+    let count = count as usize;
+    if count > configured {
+        return Err(format!(
+            "batch LoRA count {count} exceeds configured maximum {configured}"
+        ));
+    }
+    Ok(count)
+}
+
 /// Coordinator for EP workers. Lives in the HTTP server process.
 /// Holds no GPU resources — only IPC state.
 pub struct EpCoordinator {
     channel: Arc<EpChannel>,
     worker_pids: Vec<u32>,
+    dispatch_lock: Mutex<()>,
+    shutdown_started: AtomicBool,
+    transaction_sequence: AtomicU64,
 }
 
 impl EpCoordinator {
@@ -32,66 +74,504 @@ impl EpCoordinator {
     /// Each worker is pinned to GPU `rank`.
     /// Uses exec (not fork) because CUDA + fork is incompatible — forked children
     /// inherit parent's CUDA context but can't use it.
-    pub fn launch(
-        world_size: usize,
-        metrics_dir: PathBuf,
-    ) -> io::Result<Self> {
-        let channel = EpChannel::new(world_size)?;
+    pub fn launch(world_size: usize, metrics_dir: PathBuf) -> io::Result<Self> {
+        let dispatch_timeout = std::env::var("RUSTRAIN_EP_DISPATCH_TIMEOUT_SECS")
+            .ok()
+            .and_then(|value| value.parse::<u64>().ok())
+            .filter(|seconds| *seconds > 0)
+            .map(std::time::Duration::from_secs)
+            .unwrap_or_else(|| std::time::Duration::from_secs(10 * 60));
+        let channel = EpChannel::new_with_timeout(world_size, dispatch_timeout)?;
         let shm_name = channel.shm_name().to_string();
         let channel = Arc::new(channel);
 
-        let exe = std::env::current_exe()
-            .map_err(|e| io::Error::other(format!("current_exe: {e}")))?;
-        let mut worker_pids = Vec::with_capacity(world_size);
+        let exe =
+            std::env::current_exe().map_err(|e| io::Error::other(format!("current_exe: {e}")))?;
+        let mut workers = Vec::with_capacity(world_size);
 
         for rank in 0..world_size {
             let metrics_path = metrics_dir.join(format!("ep_rank{}_metrics.jsonl", rank));
             let child = std::process::Command::new(&exe)
                 .arg("ep-worker")
-                .arg("--shm-name").arg(&shm_name)
-                .arg("--rank").arg(rank.to_string())
-                .arg("--world-size").arg(world_size.to_string())
-                .arg("--metrics-path").arg(&metrics_path)
+                .arg("--shm-name")
+                .arg(&shm_name)
+                .arg("--rank")
+                .arg(rank.to_string())
+                .arg("--world-size")
+                .arg(world_size.to_string())
+                .arg("--metrics-path")
+                .arg(&metrics_path)
                 .env("RANK", rank.to_string())
                 .env("WORLD_SIZE", world_size.to_string())
                 .env("LOCAL_RANK", rank.to_string())
-                .env("QWEN36_LOSS_DIAG", std::env::var("QWEN36_LOSS_DIAG").unwrap_or_default())
-                .env("QWEN36_GROUP_SIZE", std::env::var("QWEN36_GROUP_SIZE").unwrap_or_default())
-                .env("QWEN36_FUSED_CE", std::env::var("QWEN36_FUSED_CE").unwrap_or_default())
+                .env("RUSTRAIN_NCCL_RUN_ID", &shm_name)
+                .env(
+                    "QWEN36_LOSS_DIAG",
+                    std::env::var("QWEN36_LOSS_DIAG").unwrap_or_default(),
+                )
+                .env(
+                    "QWEN36_GROUP_SIZE",
+                    std::env::var("QWEN36_GROUP_SIZE").unwrap_or_default(),
+                )
+                .env(
+                    "QWEN36_FUSED_CE",
+                    std::env::var("QWEN36_FUSED_CE").unwrap_or_default(),
+                )
                 .stdout(std::process::Stdio::inherit())
                 .stderr(std::process::Stdio::inherit())
-                .spawn()
-                .map_err(|e| io::Error::other(format!("spawn worker {rank}: {e}")))?;
+                .spawn();
+            let child = match child {
+                Ok(child) => child,
+                Err(error) => {
+                    terminate_children(&mut workers);
+                    return Err(io::Error::other(format!("spawn worker {rank}: {error}")));
+                }
+            };
             let pid = child.id();
-            // Keep child handle alive — store in a static to prevent early drop
-            // (drop would kill the child)
-            std::mem::forget(child);
-            worker_pids.push(pid);
+            workers.push(child);
             tracing::info!("Launched EP worker rank {} (PID {})", rank, pid);
         }
 
         // Wait for workers to initialize
         std::thread::sleep(std::time::Duration::from_secs(2));
+        for (rank, worker) in workers.iter_mut().enumerate() {
+            match worker.try_wait() {
+                Ok(None) => {}
+                Ok(Some(status)) => {
+                    terminate_children(&mut workers);
+                    return Err(io::Error::other(format!(
+                        "EP worker rank {rank} exited during startup with {status}"
+                    )));
+                }
+                Err(error) => {
+                    terminate_children(&mut workers);
+                    return Err(io::Error::other(format!(
+                        "inspect EP worker rank {rank} during startup: {error}"
+                    )));
+                }
+            }
+        }
+        let worker_pids = workers.iter().map(std::process::Child::id).collect();
 
-        Ok(Self { channel, worker_pids })
+        Ok(Self {
+            channel,
+            worker_pids,
+            dispatch_lock: Mutex::new(()),
+            shutdown_started: AtomicBool::new(false),
+            transaction_sequence: AtomicU64::new(0),
+        })
     }
 
     /// Dispatch a command to all workers, wait for completion, return rank 0's result.
     pub fn dispatch(&self, cmd: &EpCommand) -> EpResult {
+        let _guard = match self.dispatch_lock.lock() {
+            Ok(guard) => guard,
+            Err(error) => return EpResult::Error(format!("IPC dispatch lock poisoned: {error}")),
+        };
+        self.dispatch_locked(cmd)
+    }
+
+    pub fn dispatch_with_slab(&self, cmd: &EpCommand, payload: &[u8]) -> EpResult {
+        let _guard = match self.dispatch_lock.lock() {
+            Ok(guard) => guard,
+            Err(error) => return EpResult::Error(format!("IPC dispatch lock poisoned: {error}")),
+        };
+        if self.shutdown_started.load(Ordering::Acquire) {
+            return EpResult::Error("EP coordinator is shut down".to_string());
+        }
+        match self.channel.broadcast_with_slab(cmd, payload) {
+            Ok(result) => self.handle_dispatch_result(result),
+            Err(error) => self.handle_dispatch_error(error),
+        }
+    }
+
+    fn dispatch_locked(&self, cmd: &EpCommand) -> EpResult {
+        if self.shutdown_started.load(Ordering::Acquire) {
+            return EpResult::Error("EP coordinator is shut down".to_string());
+        }
         match self.channel.broadcast(cmd) {
+            Ok(result) => self.handle_dispatch_result(result),
+            Err(error) => self.handle_dispatch_error(error),
+        }
+    }
+
+    fn handle_dispatch_result(&self, result: EpResult) -> EpResult {
+        if is_terminal_native_context_result(&result) {
+            self.enter_terminal();
+        }
+        result
+    }
+
+    fn handle_dispatch_error(&self, error: io::Error) -> EpResult {
+        if error.kind() == io::ErrorKind::InvalidInput {
+            return EpResult::Error(format!("invalid IPC request: {error}"));
+        }
+        let exited = self.exited_workers();
+        if self
+            .shutdown_started
+            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
+            .is_ok()
+        {
+            self.terminate_workers();
+        }
+        let message = if exited.is_empty() {
+            format!("IPC error: {error}")
+        } else {
+            format!("IPC error: {error}; exited workers: {}", exited.join(", "))
+        };
+        EpResult::Error(message)
+    }
+
+    pub fn next_generation(&self, operation: &str) -> Result<String, String> {
+        let nanos = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .map_err(|error| format!("system clock precedes UNIX epoch: {error}"))?
+            .as_nanos();
+        let sequence = self.transaction_sequence.fetch_add(1, Ordering::Relaxed);
+        checkpoint_generation(std::process::id(), nanos, sequence, operation)
+    }
+
+    pub fn coordinated_save_checkpoint(
+        &self,
+        session_id: &str,
+        final_path: &str,
+    ) -> Result<(u64, f64), String> {
+        let _guard = self
+            .dispatch_lock
+            .lock()
+            .map_err(|error| format!("IPC dispatch lock poisoned: {error}"))?;
+        if self.shutdown_started.load(Ordering::Acquire) {
+            return Err("EP coordinator is shut down".into());
+        }
+
+        let generation = self.next_generation("save")?;
+        let final_path = Path::new(final_path);
+        let expected_staging = checkpoint_staging_path(final_path, &generation)?;
+        let staging_string = expected_staging
+            .to_str()
+            .ok_or_else(|| "checkpoint staging path is not UTF-8".to_string())?
+            .to_string();
+        let staging = create_checkpoint_staging(final_path, &generation)?;
+        debug_assert_eq!(staging, expected_staging);
+
+        let result = (|| {
+            let saved = match expect_checkpoint_result(
+                self.dispatch_locked(&EpCommand::PrepareSaveCheckpoint {
+                    session_id: session_id.to_string(),
+                    path: staging_string.clone(),
+                    generation: generation.clone(),
+                }),
+                "prepare distributed checkpoint save",
+            ) {
+                Ok(result) => result,
+                Err(error) => {
+                    if !self.shutdown_started.load(Ordering::Acquire) {
+                        self.abort_checkpoint_locked(session_id, &generation)?;
+                    }
+                    return Err(error);
+                }
+            };
+
+            let validated = match expect_checkpoint_result(
+                self.dispatch_locked(&EpCommand::PrepareLoadCheckpoint {
+                    session_id: session_id.to_string(),
+                    path: staging_string.clone(),
+                    transaction_id: generation.clone(),
+                }),
+                "validate distributed checkpoint staging",
+            ) {
+                Ok(result) => result,
+                Err(error) => {
+                    if !self.shutdown_started.load(Ordering::Acquire) {
+                        self.abort_checkpoint_locked(session_id, &generation)?;
+                    }
+                    return Err(error);
+                }
+            };
+            if !checkpoint_metadata_matches(saved, validated) {
+                self.abort_checkpoint_locked(session_id, &generation)?;
+                return Err(format!(
+                    "saved checkpoint metadata {saved:?} differs from validated metadata {validated:?}"
+                ));
+            }
+            self.abort_checkpoint_locked(session_id, &generation)?;
+            publish_checkpoint_noreplace(&staging, final_path)?;
+            Ok(saved)
+        })();
+
+        if result.is_err() && staging.exists() {
+            if let Err(cleanup_error) = std::fs::remove_dir_all(&staging) {
+                tracing::warn!(
+                    path = %staging.display(),
+                    error = %cleanup_error,
+                    "failed to clean checkpoint transaction staging"
+                );
+            }
+        }
+        result
+    }
+
+    pub fn coordinated_load_checkpoint(
+        &self,
+        session_id: &str,
+        path: &str,
+    ) -> Result<(u64, f64), String> {
+        let _guard = self
+            .dispatch_lock
+            .lock()
+            .map_err(|error| format!("IPC dispatch lock poisoned: {error}"))?;
+        if self.shutdown_started.load(Ordering::Acquire) {
+            return Err("EP coordinator is shut down".into());
+        }
+
+        let transaction_id = self.next_generation("load")?;
+        let prepared = match expect_checkpoint_result(
+            self.dispatch_locked(&EpCommand::PrepareLoadCheckpoint {
+                session_id: session_id.to_string(),
+                path: path.to_string(),
+                transaction_id: transaction_id.clone(),
+            }),
+            "prepare distributed checkpoint load",
+        ) {
             Ok(result) => result,
-            Err(e) => EpResult::Error(format!("IPC error: {}", e)),
+            Err(error) => {
+                if !self.shutdown_started.load(Ordering::Acquire) {
+                    self.abort_checkpoint_locked(session_id, &transaction_id)?;
+                }
+                return Err(error);
+            }
+        };
+
+        let committed = expect_checkpoint_result(
+            self.dispatch_locked(&EpCommand::CommitLoadCheckpoint {
+                session_id: session_id.to_string(),
+                transaction_id,
+            }),
+            "commit distributed checkpoint load",
+        );
+        match committed {
+            Ok(committed) if checkpoint_metadata_matches(committed, prepared) => Ok(committed),
+            Ok(committed) => {
+                self.enter_terminal();
+                Err(format!(
+                    "committed checkpoint metadata {committed:?} differs from prepared metadata {prepared:?}"
+                ))
+            }
+            Err(error) => {
+                self.enter_terminal();
+                Err(error)
+            }
+        }
+    }
+
+    fn abort_checkpoint_locked(
+        &self,
+        session_id: &str,
+        transaction_id: &str,
+    ) -> Result<(), String> {
+        match self.dispatch_locked(&EpCommand::AbortLoadCheckpoint {
+            session_id: session_id.to_string(),
+            transaction_id: transaction_id.to_string(),
+        }) {
+            EpResult::Ok => Ok(()),
+            EpResult::Error(error) => {
+                self.enter_terminal();
+                Err(format!("abort checkpoint transaction: {error}"))
+            }
+            other => {
+                self.enter_terminal();
+                Err(format!(
+                    "abort checkpoint transaction returned unexpected result: {other:?}"
+                ))
+            }
         }
     }
 
+    fn enter_terminal(&self) {
+        if self
+            .shutdown_started
+            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
+            .is_ok()
+        {
+            self.terminate_workers();
+        }
+    }
+
+    pub fn is_healthy(&self) -> bool {
+        !self.shutdown_started.load(Ordering::Acquire) && !self.channel.is_poisoned()
+    }
+
     /// Send shutdown command to all workers.
     pub fn shutdown(&self) {
-        let _ = self.channel.broadcast(&EpCommand::Shutdown);
-        // Wait for worker processes to exit
-        for pid in &self.worker_pids {
-            unsafe {
-                libc::waitpid(*pid as i32, std::ptr::null_mut(), 0);
+        if self.shutdown_started.swap(true, Ordering::AcqRel) {
+            return;
+        }
+        let _guard = self.dispatch_lock.lock().ok();
+        if !self.channel.is_poisoned() && self.channel.broadcast(&EpCommand::Shutdown).is_ok() {
+            let live = wait_for_worker_pids(&self.worker_pids, std::time::Duration::from_secs(2));
+            if !live.is_empty() {
+                terminate_worker_pids(&live);
             }
+            return;
+        }
+        self.terminate_workers();
+    }
+
+    fn exited_workers(&self) -> Vec<String> {
+        let mut exited = Vec::new();
+        for (rank, pid) in self.worker_pids.iter().enumerate() {
+            if unsafe { libc::kill(*pid as i32, 0) } != 0
+                && io::Error::last_os_error().raw_os_error() == Some(libc::ESRCH)
+            {
+                exited.push(format!("rank {rank} pid {pid}"));
+            }
+        }
+        exited
+    }
+
+    fn terminate_workers(&self) {
+        terminate_worker_pids(&self.worker_pids);
+    }
+}
+
+fn expect_checkpoint_result(result: EpResult, operation: &str) -> Result<(u64, f64), String> {
+    match result {
+        EpResult::Checkpoint { step, loss } => Ok((step, loss)),
+        EpResult::Error(error) => Err(format!("{operation}: {error}")),
+        other => Err(format!("{operation} returned unexpected result: {other:?}")),
+    }
+}
+
+fn checkpoint_metadata_matches(left: (u64, f64), right: (u64, f64)) -> bool {
+    left.0 == right.0 && left.1.to_bits() == right.1.to_bits()
+}
+
+fn checkpoint_generation(
+    pid: u32,
+    nanos: u128,
+    sequence: u64,
+    operation: &str,
+) -> Result<String, String> {
+    if operation.is_empty()
+        || !operation
+            .bytes()
+            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
+    {
+        return Err("checkpoint operation must be non-empty path-safe ASCII".into());
+    }
+    let generation = format!("ep-{pid}-{nanos}-{sequence:020}-{operation}");
+    if generation.len() > 256 {
+        return Err("checkpoint generation must be at most 256 characters".into());
+    }
+    Ok(generation)
+}
+
+fn checkpoint_staging_path(final_path: &Path, generation: &str) -> Result<PathBuf, String> {
+    if generation.is_empty()
+        || !generation
+            .bytes()
+            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
+    {
+        return Err("checkpoint generation must be non-empty path-safe ASCII".into());
+    }
+    let file_name = final_path
+        .file_name()
+        .and_then(|name| name.to_str())
+        .filter(|name| !name.is_empty() && *name != "." && *name != "..")
+        .ok_or_else(|| "checkpoint destination must have a UTF-8 file name".to_string())?;
+    Ok(final_path.with_file_name(format!(
+        ".{file_name}.rustrain-checkpoint-{generation}.partial"
+    )))
+}
+
+fn create_checkpoint_staging(final_path: &Path, generation: &str) -> Result<PathBuf, String> {
+    if final_path.exists() {
+        return Err(format!(
+            "checkpoint destination already exists: {}",
+            final_path.display()
+        ));
+    }
+    let staging = checkpoint_staging_path(final_path, generation)?;
+    std::fs::create_dir(&staging).map_err(|error| {
+        format!(
+            "create exclusive checkpoint staging {}: {error}",
+            staging.display()
+        )
+    })?;
+    Ok(staging)
+}
+
+fn publish_checkpoint_noreplace(staging: &Path, final_path: &Path) -> Result<(), String> {
+    let staging_c = CString::new(staging.as_os_str().as_bytes())
+        .map_err(|_| "checkpoint staging path contains a NUL byte".to_string())?;
+    let final_c = CString::new(final_path.as_os_str().as_bytes())
+        .map_err(|_| "checkpoint destination path contains a NUL byte".to_string())?;
+    let result = unsafe {
+        libc::renameat2(
+            libc::AT_FDCWD,
+            staging_c.as_ptr(),
+            libc::AT_FDCWD,
+            final_c.as_ptr(),
+            libc::RENAME_NOREPLACE,
+        )
+    };
+    if result == 0 {
+        Ok(())
+    } else {
+        let error = io::Error::last_os_error();
+        Err(format!(
+            "publish checkpoint {} as {} without replacement: {error}",
+            staging.display(),
+            final_path.display()
+        ))
+    }
+}
+
+fn wait_for_worker_pids(pids: &[u32], timeout: std::time::Duration) -> Vec<u32> {
+    let deadline = std::time::Instant::now() + timeout;
+    let mut live = pids.to_vec();
+    while !live.is_empty() && std::time::Instant::now() < deadline {
+        live.retain(|pid| {
+            let result = unsafe { libc::waitpid(*pid as i32, std::ptr::null_mut(), libc::WNOHANG) };
+            result == 0
+                || (result < 0 && io::Error::last_os_error().raw_os_error() == Some(libc::EINTR))
+        });
+        if !live.is_empty() {
+            std::thread::sleep(std::time::Duration::from_millis(10));
+        }
+    }
+    live
+}
+
+fn terminate_worker_pids(pids: &[u32]) {
+    for pid in pids {
+        unsafe {
+            libc::kill(*pid as i32, libc::SIGTERM);
+        }
+    }
+    let live = wait_for_worker_pids(pids, std::time::Duration::from_secs(2));
+    for pid in live {
+        unsafe {
+            libc::kill(pid as i32, libc::SIGKILL);
+        }
+        wait_for_worker_exit(pid);
+    }
+}
+
+fn terminate_children(children: &mut [std::process::Child]) {
+    for child in children.iter_mut() {
+        let _ = child.kill();
+    }
+    for child in children.iter_mut() {
+        let _ = child.wait();
+    }
+}
+
+fn wait_for_worker_exit(pid: u32) {
+    loop {
+        let result = unsafe { libc::waitpid(pid as i32, std::ptr::null_mut(), 0) };
+        if result >= 0 || io::Error::last_os_error().raw_os_error() != Some(libc::EINTR) {
+            return;
         }
     }
 }
@@ -130,7 +610,8 @@ pub fn worker_main(
         std::env::set_var("LOCAL_RANK", rank.to_string());
     }
 
-    let mut session = Qwen36Session::new(device, compute_kind, metrics_path);
+    let mut session = Qwen36Session::new(device, compute_kind, metrics_path.clone());
+    let mut active_session_id = None;
 
     // Pre-create NCCL communicator — all workers reach this point simultaneously
     // because EpCoordinator::launch forks all workers before they enter the loop.
@@ -148,7 +629,35 @@ pub fn worker_main(
             }
         };
 
-        let result = execute_command(&mut session, &cmd);
+        let result = match &cmd {
+            EpCommand::CreateSession { session_id } => {
+                match create_worker_session(&mut active_session_id, session_id) {
+                    Ok(()) => EpResult::Ok,
+                    Err(error) => EpResult::Error(error),
+                }
+            }
+            EpCommand::DeleteSession { session_id } => {
+                match delete_worker_session(&mut active_session_id, session_id) {
+                    Ok(()) => {
+                        session = Qwen36Session::new(device, compute_kind, metrics_path.clone());
+                        EpResult::Ok
+                    }
+                    Err(error) => EpResult::Error(error),
+                }
+            }
+            EpCommand::Shutdown => execute_command(&mut session, &worker, &cmd),
+            _ => match require_worker_session(&active_session_id, command_session_id(&cmd)) {
+                Ok(()) => execute_command(&mut session, &worker, &cmd),
+                Err(error) => EpResult::Error(error),
+            },
+        };
+        let result = if session.native_context_is_healthy() {
+            result
+        } else {
+            EpResult::Error(format!(
+                "{TERMINAL_NATIVE_CONTEXT_PREFIX}: worker group must be recreated"
+            ))
+        };
 
         if let Err(e) = worker.signal_done(&result) {
             eprintln!("[ep-worker-{}] signal_done error: {}", rank, e);
@@ -165,18 +674,16 @@ pub fn worker_main(
 }
 
 /// Execute a command on the local session, return result.
-fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult {
+fn execute_command(session: &mut Qwen36Session, worker: &EpWorker, cmd: &EpCommand) -> EpResult {
     match cmd {
-        EpCommand::CreateSession { session_id: _ } => {
-            // Session already created in worker_main; just acknowledge
-            EpResult::Ok
-        }
-        EpCommand::DeleteSession { .. } => {
-            // In EP mode, we don't really delete the session — just reset state
-            // A full implementation would manage session lifecycle per worker
-            EpResult::Ok
+        EpCommand::CreateSession { .. } | EpCommand::DeleteSession { .. } => {
+            EpResult::Error("session lifecycle command reached the execution layer".into())
         }
-        EpCommand::LoadModel { model_path, config_toml, .. } => {
+        EpCommand::LoadModel {
+            model_path,
+            config_toml,
+            ..
+        } => {
             match session.load_model(SessLoadModelRequest {
                 model_path: model_path.clone(),
                 config_toml: config_toml.clone(),
@@ -185,7 +692,11 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult {
                 Err(e) => EpResult::Error(e.to_string()),
             }
         }
-        EpCommand::LoadDataset { jsonl_path, seq_len, .. } => {
+        EpCommand::LoadDataset {
+            jsonl_path,
+            seq_len,
+            ..
+        } => {
             match session.load_dataset(SessLoadDatasetRequest {
                 jsonl_path: jsonl_path.clone(),
                 seq_len: *seq_len,
@@ -194,7 +705,17 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult {
                 Err(e) => EpResult::Error(e.to_string()),
             }
         }
-        EpCommand::InitLora { rank, alpha, target_layers, target_modules, lr, beta1, beta2, eps, .. } => {
+        EpCommand::InitLora {
+            rank,
+            alpha,
+            target_layers,
+            target_modules,
+            lr,
+            beta1,
+            beta2,
+            eps,
+            ..
+        } => {
             match session.init_lora(InitLoRARequest {
                 rank: *rank,
                 alpha: *alpha,
@@ -209,92 +730,441 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult {
                 Err(e) => EpResult::Error(e.to_string()),
             }
         }
-        EpCommand::AddLora { rank, alpha, target_layers, target_modules, .. } => {
+        EpCommand::AddLora {
+            rank,
+            alpha,
+            target_layers,
+            target_modules,
+            optimizer_lr,
+            optimizer_beta1,
+            optimizer_beta2,
+            optimizer_eps,
+            ..
+        } => {
             match session.add_lora(AddLoRARequest {
                 rank: *rank,
                 alpha: *alpha,
                 target_layers: target_layers.clone(),
                 target_modules: target_modules.clone(),
+                optimizer_lr: *optimizer_lr,
+                optimizer_beta1: *optimizer_beta1,
+                optimizer_beta2: *optimizer_beta2,
+                optimizer_eps: *optimizer_eps,
             }) {
                 Ok(id) => EpResult::AdapterId(id),
                 Err(e) => EpResult::Error(e.to_string()),
             }
         }
-        EpCommand::BatchAddLora { count, rank, alpha, target_layers, target_modules, .. } => {
-            let mut ids = Vec::with_capacity(*count as usize);
-            for _ in 0..*count {
+        EpCommand::BatchAddLora {
+            count,
+            rank,
+            alpha,
+            target_layers,
+            target_modules,
+            optimizer_lr,
+            optimizer_beta1,
+            optimizer_beta2,
+            optimizer_eps,
+            ..
+        } => {
+            let count = match validate_batch_add_lora_count(*count) {
+                Ok(count) => count,
+                Err(error) => return EpResult::Error(error),
+            };
+            let mut ids = Vec::with_capacity(count);
+            for _ in 0..count {
                 match session.add_lora(AddLoRARequest {
                     rank: *rank,
                     alpha: *alpha,
                     target_layers: target_layers.clone(),
                     target_modules: target_modules.clone(),
+                    optimizer_lr: *optimizer_lr,
+                    optimizer_beta1: *optimizer_beta1,
+                    optimizer_beta2: *optimizer_beta2,
+                    optimizer_eps: *optimizer_eps,
                 }) {
                     Ok(id) => ids.push(id),
-                    Err(e) => { return EpResult::Error(e.to_string()); }
+                    Err(e) => {
+                        let error = e.to_string();
+                        let mut rollback_errors = Vec::new();
+                        for id in ids.into_iter().rev() {
+                            if let Err(rollback_error) = session.remove_lora(id) {
+                                rollback_errors.push(format!("adapter {id}: {rollback_error}"));
+                            }
+                        }
+                        if rollback_errors.is_empty() {
+                            return EpResult::Error(error);
+                        }
+                        return EpResult::Error(format!(
+                            "{error}; batch rollback failed: {}",
+                            rollback_errors.join("; ")
+                        ));
+                    }
                 }
             }
             EpResult::Count(ids.len() as usize)
         }
-        EpCommand::RemoveLora { adapter_id, .. } => {
-            match session.remove_lora(*adapter_id) {
-                Ok(b) => {
-                    if b { EpResult::Ok } else { EpResult::Error("adapter not found".into()) }
+        EpCommand::RemoveLora { adapter_id, .. } => match session.remove_lora(*adapter_id) {
+            Ok(b) => {
+                if b {
+                    EpResult::Ok
+                } else {
+                    EpResult::Error("adapter not found".into())
                 }
+            }
+            Err(e) => EpResult::Error(e.to_string()),
+        },
+        EpCommand::ListLora { .. } => EpResult::AdapterIds(session.list_lora()),
+        EpCommand::TrainStep {
+            input_ids,
+            target_mask,
+            attention_mask,
+            batch_size,
+            seq_len,
+            ..
+        } => {
+            let source_shard = match source_shard_from_env() {
+                Ok(shard) => shard,
+                Err(error) => return EpResult::Error(error),
+            };
+            let rows = match train_step_rows(*batch_size, source_shard) {
+                Ok(rows) => rows,
+                Err(error) => return EpResult::Error(error),
+            };
+            let elements = match tensor_element_range(*batch_size, *seq_len, rows.clone()) {
+                Ok(elements) => elements,
+                Err(error) => return EpResult::Error(error),
+            };
+            if let Err(error) = validate_flat_tensor_lengths(
+                *batch_size,
+                *seq_len,
+                input_ids.len(),
+                target_mask.len(),
+                attention_mask.len(),
+            ) {
+                return EpResult::Error(error);
+            }
+            match session.train_step_host_i64(
+                &input_ids[elements.clone()],
+                &target_mask[elements.clone()],
+                &attention_mask[elements],
+                rows.len(),
+                *seq_len,
+            ) {
+                Ok(TrainOutput { loss, step }) => EpResult::Train { loss, step },
                 Err(e) => EpResult::Error(e.to_string()),
             }
         }
-        EpCommand::ListLora { .. } => {
-            EpResult::AdapterIds(session.list_lora())
+        EpCommand::TrainStepSlab { tensors, .. } => {
+            let (input_ids, target_mask, attention_mask) = match slab_tensor_views(worker, tensors)
+            {
+                Ok(views) => views,
+                Err(error) => return EpResult::Error(error),
+            };
+            let source_shard = match source_shard_from_env() {
+                Ok(shard) => shard,
+                Err(error) => return EpResult::Error(error),
+            };
+            let rows = match train_step_rows(tensors.batch_size, source_shard) {
+                Ok(rows) => rows,
+                Err(error) => return EpResult::Error(error),
+            };
+            let elements =
+                match tensor_element_range(tensors.batch_size, tensors.seq_len, rows.clone()) {
+                    Ok(elements) => elements,
+                    Err(error) => return EpResult::Error(error),
+                };
+            match session.train_step_host_i64(
+                &input_ids[elements.clone()],
+                &target_mask[elements.clone()],
+                &attention_mask[elements],
+                rows.len(),
+                tensors.seq_len,
+            ) {
+                Ok(TrainOutput { loss, step }) => EpResult::Train { loss, step },
+                Err(error) => EpResult::Error(error.to_string()),
+            }
         }
-        EpCommand::TrainStep { input_ids, target_mask, attention_mask, seq_len, .. } => {
-            let sl = *seq_len as i64;
-            let input_ids_tensor = tch::Tensor::from_slice(input_ids).reshape(&[1, sl]).to_device(session.device());
-            let target_mask_tensor = tch::Tensor::from_slice(target_mask).reshape(&[1, sl]).to_device(session.device());
-            let attention_mask_tensor = tch::Tensor::from_slice(attention_mask).reshape(&[1, sl]).to_device(session.device());
-
-            match session.train_step(TrainInput {
-                input_ids: input_ids_tensor,
-                target_mask: target_mask_tensor,
-                attention_mask: attention_mask_tensor,
-            }) {
-                Ok(TrainOutput { loss, .. }) => EpResult::Loss(loss),
-                Err(e) => EpResult::Error(e.to_string()),
+        EpCommand::TrainMultiLora {
+            input_ids,
+            target_mask,
+            attention_mask,
+            batch_size,
+            seq_len,
+            n_total,
+            lora_rank,
+            adapter_ids,
+            expected_steps,
+            ..
+        } => {
+            if let Err(error) = validate_flat_tensor_lengths(
+                *batch_size,
+                *seq_len,
+                input_ids.len(),
+                target_mask.len(),
+                attention_mask.len(),
+            ) {
+                return EpResult::Error(error);
+            }
+            if let Err(error) = session.validate_dynamic_adapter_steps(adapter_ids, expected_steps)
+            {
+                return EpResult::Error(error.to_string());
+            }
+            let source_shard = match source_shard_from_env() {
+                Ok(shard) => shard,
+                Err(error) => return EpResult::Error(error),
+            };
+            let rows = match multi_lora_rows(*batch_size, *n_total, source_shard) {
+                Ok(rows) => rows,
+                Err(error) => return EpResult::Error(error),
+            };
+            let elements = match tensor_element_range(*batch_size, *seq_len, rows.clone()) {
+                Ok(elements) => elements,
+                Err(error) => return EpResult::Error(error),
+            };
+            if adapter_ids.is_empty() {
+                match session.train_multi_lora_host_i64(
+                    &input_ids[elements.clone()],
+                    &target_mask[elements.clone()],
+                    &attention_mask[elements],
+                    rows.len(),
+                    *seq_len,
+                    *n_total,
+                    *lora_rank,
+                    adapter_ids,
+                ) {
+                    Ok(TrainOutput { loss, step }) => EpResult::Train { loss, step },
+                    Err(error) => EpResult::Error(error.to_string()),
+                }
+            } else {
+                match session.train_multi_lora_host_i64_report(
+                    &input_ids[elements.clone()],
+                    &target_mask[elements.clone()],
+                    &attention_mask[elements],
+                    rows.len(),
+                    *seq_len,
+                    *n_total,
+                    *lora_rank,
+                    adapter_ids,
+                ) {
+                    Ok(output) => EpResult::MultiLoraTrain {
+                        loss: output.loss,
+                        step: output.step,
+                        adapter_losses: adapter_ids
+                            .iter()
+                            .copied()
+                            .zip(output.adapter_losses)
+                            .map(|(adapter_id, loss)| rustrain_ipc::command::AdapterLoss {
+                                adapter_id,
+                                loss,
+                            })
+                            .collect(),
+                        adapter_steps: adapter_ids
+                            .iter()
+                            .copied()
+                            .zip(output.adapter_steps)
+                            .map(|(adapter_id, step)| rustrain_ipc::command::AdapterStep {
+                                adapter_id,
+                                step,
+                            })
+                            .collect(),
+                    },
+                    Err(error) => EpResult::Error(error.to_string()),
+                }
             }
         }
-        EpCommand::TrainMultiLora { input_ids, target_mask, attention_mask, seq_len, n_total, lora_rank, .. } => {
-            let sl = *seq_len as i64;
-            let input_ids_tensor = tch::Tensor::from_slice(input_ids).reshape(&[1, sl]).to_device(session.device());
-            let target_mask_tensor = tch::Tensor::from_slice(target_mask).reshape(&[1, sl]).to_device(session.device());
-            let attention_mask_tensor = tch::Tensor::from_slice(attention_mask).reshape(&[1, sl]).to_device(session.device());
-
-            match session.train_multi_lora(TrainInput {
-                input_ids: input_ids_tensor,
-                target_mask: target_mask_tensor,
-                attention_mask: attention_mask_tensor,
-            }, *n_total, *lora_rank) {
-                Ok(TrainOutput { loss, .. }) => EpResult::Loss(loss),
-                Err(e) => EpResult::Error(e.to_string()),
+        EpCommand::TrainMultiLoraSlab {
+            tensors,
+            n_total,
+            lora_rank,
+            adapter_ids,
+            expected_steps,
+            ..
+        } => {
+            let (input_ids, target_mask, attention_mask) = match slab_tensor_views(worker, tensors)
+            {
+                Ok(views) => views,
+                Err(error) => return EpResult::Error(error),
+            };
+            if let Err(error) = session.validate_dynamic_adapter_steps(adapter_ids, expected_steps)
+            {
+                return EpResult::Error(error.to_string());
+            }
+            let source_shard = match source_shard_from_env() {
+                Ok(shard) => shard,
+                Err(error) => return EpResult::Error(error),
+            };
+            let rows = match multi_lora_rows(tensors.batch_size, *n_total, source_shard) {
+                Ok(rows) => rows,
+                Err(error) => return EpResult::Error(error),
+            };
+            let elements =
+                match tensor_element_range(tensors.batch_size, tensors.seq_len, rows.clone()) {
+                    Ok(elements) => elements,
+                    Err(error) => return EpResult::Error(error),
+                };
+            if adapter_ids.is_empty() {
+                match session.train_multi_lora_host_i64(
+                    &input_ids[elements.clone()],
+                    &target_mask[elements.clone()],
+                    &attention_mask[elements],
+                    rows.len(),
+                    tensors.seq_len,
+                    *n_total,
+                    *lora_rank,
+                    adapter_ids,
+                ) {
+                    Ok(TrainOutput { loss, step }) => EpResult::Train { loss, step },
+                    Err(error) => EpResult::Error(error.to_string()),
+                }
+            } else {
+                match session.train_multi_lora_host_i64_report(
+                    &input_ids[elements.clone()],
+                    &target_mask[elements.clone()],
+                    &attention_mask[elements],
+                    rows.len(),
+                    tensors.seq_len,
+                    *n_total,
+                    *lora_rank,
+                    adapter_ids,
+                ) {
+                    Ok(output) => EpResult::MultiLoraTrain {
+                        loss: output.loss,
+                        step: output.step,
+                        adapter_losses: adapter_ids
+                            .iter()
+                            .copied()
+                            .zip(output.adapter_losses)
+                            .map(|(adapter_id, loss)| rustrain_ipc::command::AdapterLoss {
+                                adapter_id,
+                                loss,
+                            })
+                            .collect(),
+                        adapter_steps: adapter_ids
+                            .iter()
+                            .copied()
+                            .zip(output.adapter_steps)
+                            .map(|(adapter_id, step)| rustrain_ipc::command::AdapterStep {
+                                adapter_id,
+                                step,
+                            })
+                            .collect(),
+                    },
+                    Err(error) => EpResult::Error(error.to_string()),
+                }
             }
         }
-        EpCommand::EvalStep { input_ids, target_mask, attention_mask, seq_len, .. } => {
-            let sl = *seq_len as i64;
-            let input_ids_tensor = tch::Tensor::from_slice(input_ids).reshape(&[1, sl]).to_device(session.device());
-            let target_mask_tensor = tch::Tensor::from_slice(target_mask).reshape(&[1, sl]).to_device(session.device());
-            let attention_mask_tensor = tch::Tensor::from_slice(attention_mask).reshape(&[1, sl]).to_device(session.device());
-
-            match session.eval_step(TrainInput {
-                input_ids: input_ids_tensor,
-                target_mask: target_mask_tensor,
-                attention_mask: attention_mask_tensor,
-            }) {
+        EpCommand::EvalStep {
+            input_ids,
+            target_mask,
+            attention_mask,
+            seq_len,
+            ..
+        } => {
+            if let Err(error) = validate_flat_tensor_lengths(
+                1,
+                *seq_len,
+                input_ids.len(),
+                target_mask.len(),
+                attention_mask.len(),
+            ) {
+                return EpResult::Error(error);
+            }
+            match session.eval_step_host_i64(input_ids, target_mask, attention_mask, 1, *seq_len) {
                 Ok(EvalOutput { loss }) => EpResult::Loss(loss),
                 Err(e) => EpResult::Error(e.to_string()),
             }
         }
-        EpCommand::ExportAdapter { path, .. } => {
-            match session.export_adapter(path) {
-                Ok(n) => EpResult::Count(n),
-                Err(e) => EpResult::Error(e.to_string()),
+        EpCommand::EvalStepSlab { tensors, .. } => {
+            let (input_ids, target_mask, attention_mask) = match slab_tensor_views(worker, tensors)
+            {
+                Ok(views) => views,
+                Err(error) => return EpResult::Error(error),
+            };
+            match session.eval_step_host_i64(
+                input_ids,
+                target_mask,
+                attention_mask,
+                tensors.batch_size,
+                tensors.seq_len,
+            ) {
+                Ok(EvalOutput { loss }) => EpResult::Loss(loss),
+                Err(error) => EpResult::Error(error.to_string()),
+            }
+        }
+        EpCommand::EvalMultiLoraSlab {
+            tensors,
+            adapter_ids,
+            ..
+        } => {
+            if adapter_ids.is_empty() || adapter_ids.iter().any(|id| *id <= 0) {
+                return EpResult::Error("adapter_ids must contain positive IDs".into());
+            }
+            if tensors.batch_size != 1 && tensors.batch_size != adapter_ids.len() {
+                return EpResult::Error(format!(
+                    "selected eval batch_size must be 1 or adapter count {}, got {}",
+                    adapter_ids.len(),
+                    tensors.batch_size
+                ));
+            }
+            let (input_ids, target_mask, attention_mask) = match slab_tensor_views(worker, tensors)
+            {
+                Ok(views) => views,
+                Err(error) => return EpResult::Error(error),
+            };
+            match session.eval_multi_lora_host_i64(
+                input_ids,
+                target_mask,
+                attention_mask,
+                tensors.batch_size,
+                tensors.seq_len,
+                adapter_ids,
+            ) {
+                Ok(MultiLoraEvalOutput { adapter_losses }) => EpResult::MultiLoraEval {
+                    adapter_losses: adapter_losses
+                        .into_iter()
+                        .map(|(adapter_id, loss)| rustrain_ipc::command::AdapterLoss {
+                            adapter_id,
+                            loss,
+                        })
+                        .collect(),
+                },
+                Err(error) => EpResult::Error(error.to_string()),
+            }
+        }
+        EpCommand::ExportAdapter {
+            path,
+            adapter_id,
+            generation,
+            ..
+        } => match session.export_distributed_adapter(path, *adapter_id, generation) {
+            Ok(n) => EpResult::Count(n),
+            Err(e) => EpResult::Error(e.to_string()),
+        },
+        EpCommand::PrepareSaveCheckpoint {
+            path, generation, ..
+        } => match session.save_checkpoint_with_generation(path, Some(generation)) {
+            Ok((step, loss)) => EpResult::Checkpoint { step, loss },
+            Err(error) => EpResult::Error(error.to_string()),
+        },
+        EpCommand::PrepareLoadCheckpoint {
+            path,
+            transaction_id,
+            ..
+        } => match session.prepare_checkpoint_load(path, transaction_id) {
+            Ok((step, loss)) => EpResult::Checkpoint { step, loss },
+            Err(error) => EpResult::Error(error.to_string()),
+        },
+        EpCommand::CommitLoadCheckpoint { transaction_id, .. } => {
+            match session.commit_checkpoint_load(transaction_id) {
+                Ok((step, loss)) => EpResult::Checkpoint { step, loss },
+                Err(error) => EpResult::Error(error.to_string()),
+            }
+        }
+        EpCommand::AbortLoadCheckpoint { transaction_id, .. } => {
+            match session.abort_checkpoint_load(transaction_id) {
+                Ok(()) => EpResult::Ok,
+                Err(error) => EpResult::Error(error.to_string()),
             }
         }
         EpCommand::Status { .. } => {
@@ -306,8 +1176,538 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult {
                 model_path: s.model_path,
             }
         }
-        EpCommand::Shutdown => {
-            EpResult::Ok
+        EpCommand::Shutdown => EpResult::Ok,
+    }
+}
+
+fn command_session_id(cmd: &EpCommand) -> &str {
+    match cmd {
+        EpCommand::CreateSession { session_id }
+        | EpCommand::DeleteSession { session_id }
+        | EpCommand::LoadModel { session_id, .. }
+        | EpCommand::LoadDataset { session_id, .. }
+        | EpCommand::InitLora { session_id, .. }
+        | EpCommand::AddLora { session_id, .. }
+        | EpCommand::BatchAddLora { session_id, .. }
+        | EpCommand::RemoveLora { session_id, .. }
+        | EpCommand::ListLora { session_id }
+        | EpCommand::TrainStep { session_id, .. }
+        | EpCommand::TrainStepSlab { session_id, .. }
+        | EpCommand::TrainMultiLora { session_id, .. }
+        | EpCommand::TrainMultiLoraSlab { session_id, .. }
+        | EpCommand::EvalStep { session_id, .. }
+        | EpCommand::EvalStepSlab { session_id, .. }
+        | EpCommand::EvalMultiLoraSlab { session_id, .. }
+        | EpCommand::ExportAdapter { session_id, .. }
+        | EpCommand::PrepareSaveCheckpoint { session_id, .. }
+        | EpCommand::PrepareLoadCheckpoint { session_id, .. }
+        | EpCommand::CommitLoadCheckpoint { session_id, .. }
+        | EpCommand::AbortLoadCheckpoint { session_id, .. }
+        | EpCommand::Status { session_id } => session_id,
+        EpCommand::Shutdown => "",
+    }
+}
+
+fn create_worker_session(active: &mut Option<String>, requested: &str) -> Result<(), String> {
+    if requested.is_empty() {
+        return Err("session_id must be non-empty".into());
+    }
+    match active {
+        Some(current) if current == requested => Ok(()),
+        Some(current) => Err(format!(
+            "distributed worker group already owns session {current}; use dynamic LoRA adapters for multi-tenant training or delete it before creating {requested}"
+        )),
+        None => {
+            *active = Some(requested.to_string());
+            Ok(())
+        }
+    }
+}
+
+fn require_worker_session(active: &Option<String>, requested: &str) -> Result<(), String> {
+    match active {
+        Some(current) if current == requested => Ok(()),
+        Some(current) => Err(format!(
+            "command targets session {requested}, but this distributed worker group owns {current}"
+        )),
+        None => Err(format!(
+            "command targets session {requested}, but no distributed session has been created"
+        )),
+    }
+}
+
+fn delete_worker_session(active: &mut Option<String>, requested: &str) -> Result<(), String> {
+    require_worker_session(active, requested)?;
+    *active = None;
+    Ok(())
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+struct SourceShard {
+    rank: usize,
+    size: usize,
+}
+
+fn source_shard_from_env() -> Result<Option<SourceShard>, String> {
+    let topology = ParallelTopology::from_env()
+        .map_err(|error| format!("invalid source-parallel topology: {error}"))?;
+    let global_rank = std::env::var("RANK")
+        .map_err(|_| "RANK is required for source-parallel training".to_string())?
+        .parse::<usize>()
+        .map_err(|_| "RANK must be a non-negative integer".to_string())?;
+    let ep_source_sharded = std::env::var("QWEN36_EP_A2A_SHARDED")
+        .map(|value| !value.is_empty() && value != "0")
+        .unwrap_or(false);
+    source_shard_for_topology(&topology, global_rank, ep_source_sharded)
+}
+
+pub(crate) fn validate_multi_lora_global_batch_size(
+    batch_size: usize,
+    n_total: i32,
+    world_size: usize,
+) -> Result<(), String> {
+    let topology = ParallelTopology::from_env_with_world_size(world_size)
+        .map_err(|error| format!("invalid source-parallel topology: {error}"))?;
+    let ep_source_sharded = std::env::var("QWEN36_EP_A2A_SHARDED")
+        .map(|value| !value.is_empty() && value != "0")
+        .unwrap_or(false);
+    // Admission needs only the global source-parallel size. Rank zero is a
+    // valid representative and avoids requiring worker-only RANK in the API.
+    let source_shard = source_shard_for_topology(&topology, 0, ep_source_sharded)?;
+    multi_lora_rows(batch_size, n_total, source_shard).map(|_| ())
+}
+
+fn source_shard_for_topology(
+    topology: &ParallelTopology,
+    global_rank: usize,
+    ep_source_sharded: bool,
+) -> Result<Option<SourceShard>, String> {
+    let dp_size = topology.data_parallel_size();
+    let ep_size = topology.expert_model_parallel_size();
+    let size = if ep_source_sharded {
+        dp_size
+            .checked_mul(ep_size)
+            .ok_or_else(|| "source-parallel size overflowed usize".to_string())?
+    } else {
+        dp_size
+    };
+    if size <= 1 {
+        return Ok(None);
+    }
+    let dp_rank = topology
+        .data_rank(global_rank)
+        .map_err(|error| format!("invalid source-parallel DP rank: {error}"))?;
+    let rank = if ep_source_sharded {
+        let ep_rank = topology
+            .expert_rank(global_rank)
+            .map_err(|error| format!("invalid source-parallel EP rank: {error}"))?;
+        dp_rank
+            .checked_mul(ep_size)
+            .and_then(|base| base.checked_add(ep_rank))
+            .ok_or_else(|| "source-parallel rank overflowed usize".to_string())?
+    } else {
+        dp_rank
+    };
+    Ok(Some(SourceShard { rank, size }))
+}
+
+fn validate_flat_tensor_lengths(
+    batch_size: usize,
+    seq_len: usize,
+    input_len: usize,
+    target_len: usize,
+    attention_len: usize,
+) -> Result<(), String> {
+    if batch_size == 0 || seq_len == 0 {
+        return Err(format!(
+            "batch_size and seq_len must be positive, got batch_size={batch_size} seq_len={seq_len}"
+        ));
+    }
+    let expected = batch_size
+        .checked_mul(seq_len)
+        .ok_or_else(|| "batch tensor element count overflowed usize".to_string())?;
+    if input_len != expected || target_len != expected || attention_len != expected {
+        return Err(format!(
+            "tensor lengths must match batch_size={batch_size} seq_len={seq_len} (expected {expected}), got input={input_len} target={target_len} attention={attention_len}"
+        ));
+    }
+    Ok(())
+}
+
+fn slab_tensor_views<'a>(
+    worker: &'a EpWorker,
+    tensors: &TensorSlabRef,
+) -> Result<(&'a [i64], &'a [i64], &'a [i64]), String> {
+    // The worker signals completion only after execute_command returns, so
+    // these slab borrows cannot outlive the current broadcast.
+    unsafe {
+        Ok((
+            worker
+                .slab_i64(tensors.input_ids)
+                .map_err(|error| error.to_string())?,
+            worker
+                .slab_i64(tensors.target_mask)
+                .map_err(|error| error.to_string())?,
+            worker
+                .slab_i64(tensors.attention_mask)
+                .map_err(|error| error.to_string())?,
+        ))
+    }
+}
+
+fn train_step_rows(batch_size: usize, shard: Option<SourceShard>) -> Result<Range<usize>, String> {
+    let Some(shard) = shard else {
+        return Ok(0..batch_size);
+    };
+    if shard.size == 0 || shard.rank >= shard.size {
+        return Err(format!(
+            "invalid source shard rank={}/size={}",
+            shard.rank, shard.size
+        ));
+    }
+    if batch_size % shard.size != 0 {
+        return Err(format!(
+            "global batch_size={batch_size} must be divisible by source-parallel size {}",
+            shard.size
+        ));
+    }
+    let local_batch = batch_size / shard.size;
+    let start = shard.rank * local_batch;
+    Ok(start..start + local_batch)
+}
+
+fn multi_lora_rows(
+    batch_size: usize,
+    n_total: i32,
+    shard: Option<SourceShard>,
+) -> Result<Range<usize>, String> {
+    let pp_size = std::env::var("PP_SIZE")
+        .or_else(|_| std::env::var("RUSTRAIN_PP_SIZE"))
+        .or_else(|_| std::env::var("PIPELINE_MODEL_PARALLEL_SIZE"))
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(1);
+    let pipeline_microbatches = if pp_size > 1 {
+        configured_dynamic_pipeline_microbatches().map_err(|error| error.to_string())?
+    } else {
+        1
+    };
+    multi_lora_rows_with_microbatches(batch_size, n_total, shard, pipeline_microbatches)
+}
+
+fn multi_lora_rows_with_microbatches(
+    batch_size: usize,
+    n_total: i32,
+    shard: Option<SourceShard>,
+    pipeline_microbatches: usize,
+) -> Result<Range<usize>, String> {
+    if n_total <= 0 {
+        return Err(format!("n_total must be positive, got {n_total}"));
+    }
+    if pipeline_microbatches == 0 {
+        return Err("pipeline_microbatches must be positive".to_string());
+    }
+    let n_total = n_total as usize;
+    let Some(shard) = shard else {
+        let expected_rows = n_total
+            .checked_mul(pipeline_microbatches)
+            .ok_or_else(|| "multi-LoRA row count overflowed usize".to_string())?;
+        if batch_size == 1 || batch_size == expected_rows {
+            return Ok(0..batch_size);
+        }
+        return Err(format!(
+            "multi-LoRA batch_size={batch_size} must be 1 or n_total*microbatches={} when source parallelism is disabled",
+            expected_rows
+        ));
+    };
+    if shard.size == 0 || shard.rank >= shard.size {
+        return Err(format!(
+            "invalid source shard rank={}/size={}",
+            shard.rank, shard.size
+        ));
+    }
+    let global_rows = n_total
+        .checked_mul(shard.size)
+        .and_then(|rows| rows.checked_mul(pipeline_microbatches))
+        .ok_or_else(|| "multi-LoRA global row count overflowed usize".to_string())?;
+    if batch_size != global_rows {
+        return Err(format!(
+            "source-parallel multi-LoRA batch_size={batch_size} must equal n_total*source_parallel_size*microbatches={global_rows}; submit the complete global source batch"
+        ));
+    }
+    let local_rows = n_total
+        .checked_mul(pipeline_microbatches)
+        .ok_or_else(|| "multi-LoRA local row count overflowed usize".to_string())?;
+    let start = shard
+        .rank
+        .checked_mul(local_rows)
+        .ok_or_else(|| "multi-LoRA source row start overflowed usize".to_string())?;
+    Ok(start..start + local_rows)
+}
+
+fn tensor_element_range(
+    batch_size: usize,
+    seq_len: usize,
+    rows: Range<usize>,
+) -> Result<Range<usize>, String> {
+    if rows.start > rows.end || rows.end > batch_size {
+        return Err(format!(
+            "row range {:?} is outside batch_size={batch_size}",
+            rows
+        ));
+    }
+    let start = rows
+        .start
+        .checked_mul(seq_len)
+        .ok_or_else(|| "tensor slice start overflowed usize".to_string())?;
+    let end = rows
+        .end
+        .checked_mul(seq_len)
+        .ok_or_else(|| "tensor slice end overflowed usize".to_string())?;
+    Ok(start..end)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{
+        checkpoint_generation, checkpoint_metadata_matches, checkpoint_staging_path,
+        create_checkpoint_staging, create_worker_session, delete_worker_session,
+        is_terminal_native_context_result, multi_lora_rows, multi_lora_rows_with_microbatches,
+        publish_checkpoint_noreplace,
+        require_worker_session, source_shard_for_topology, tensor_element_range, train_step_rows,
+        validate_batch_add_lora_count, validate_flat_tensor_lengths,
+        validate_multi_lora_global_batch_size, SourceShard,
+    };
+    use rustrain_ipc::EpResult;
+    use rustrain_parallel::topology::ParallelTopology;
+
+    #[test]
+    fn train_step_slices_global_batch_contiguously_by_source_rank() {
+        assert_eq!(
+            train_step_rows(8, Some(SourceShard { rank: 0, size: 2 })).unwrap(),
+            0..4
+        );
+        assert_eq!(
+            train_step_rows(8, Some(SourceShard { rank: 1, size: 2 })).unwrap(),
+            4..8
+        );
+        assert!(train_step_rows(7, Some(SourceShard { rank: 0, size: 2 })).is_err());
+        assert_eq!(train_step_rows(7, None).unwrap(), 0..7);
+    }
+
+    #[test]
+    fn batch_add_lora_count_is_fail_closed_before_allocation() {
+        assert_eq!(validate_batch_add_lora_count(1).unwrap(), 1);
+        assert_eq!(validate_batch_add_lora_count(64).unwrap(), 64);
+        assert!(validate_batch_add_lora_count(0).is_err());
+        assert!(validate_batch_add_lora_count(-1).is_err());
+        assert!(validate_batch_add_lora_count(4097).is_err());
+    }
+
+    #[test]
+    fn native_context_poison_is_a_terminal_worker_result() {
+        assert!(is_terminal_native_context_result(&EpResult::Error(
+            "terminal native LoRA context: worker group must be recreated".into()
+        )));
+        assert!(!is_terminal_native_context_result(&EpResult::Error(
+            "ordinary request validation failed".into()
+        )));
+    }
+
+    #[test]
+    fn tp_peers_with_the_same_dp_rank_select_the_same_rows() {
+        let topology = ParallelTopology::new(2, 1, 2, 1, 1).unwrap();
+        let tp_rank_zero_rows =
+            train_step_rows(12, source_shard_for_topology(&topology, 2, false).unwrap()).unwrap();
+        let tp_rank_one_rows =
+            train_step_rows(12, source_shard_for_topology(&topology, 3, false).unwrap()).unwrap();
+        assert_eq!(tp_rank_zero_rows, tp_rank_one_rows);
+        assert_eq!(tp_rank_zero_rows, 6..12);
+    }
+
+    #[test]
+    fn custom_rank_order_keeps_tp_peers_on_the_same_dp_shard() {
+        let topology = ParallelTopology::with_order(2, 1, 2, 1, 1, "dp-tp").unwrap();
+        let first_tp_peer =
+            train_step_rows(12, source_shard_for_topology(&topology, 1, false).unwrap()).unwrap();
+        let second_tp_peer =
+            train_step_rows(12, source_shard_for_topology(&topology, 3, false).unwrap()).unwrap();
+        assert_eq!(first_tp_peer, second_tp_peer);
+        assert_eq!(first_tp_peer, 6..12);
+    }
+
+    #[test]
+    fn tp_ep_dp_source_shards_follow_ep_policy() {
+        let topology = ParallelTopology::new(2, 1, 2, 2, 1).unwrap();
+        for global_rank in 0..topology.world_size() {
+            let coordinates = topology.coordinates(global_rank).unwrap();
+            let sharded = source_shard_for_topology(&topology, global_rank, true)
+                .unwrap()
+                .unwrap();
+            assert_eq!(sharded.size, 4);
+            assert_eq!(sharded.rank, coordinates.data * 2 + coordinates.expert);
+
+            let replicated = source_shard_for_topology(&topology, global_rank, false)
+                .unwrap()
+                .unwrap();
+            assert_eq!(replicated.size, 2);
+            assert_eq!(replicated.rank, coordinates.data);
+        }
+    }
+
+    #[test]
+    fn tp_peers_share_tri_axis_source_rows() {
+        let topology = ParallelTopology::with_order(2, 1, 2, 2, 1, "ep-dp-tp").unwrap();
+        for dp_rank in 0..2 {
+            for ep_rank in 0..2 {
+                let peers = (0..topology.world_size())
+                    .filter(|global_rank| {
+                        let coordinates = topology.coordinates(*global_rank).unwrap();
+                        coordinates.data == dp_rank && coordinates.expert == ep_rank
+                    })
+                    .collect::<Vec<_>>();
+                assert_eq!(peers.len(), 2);
+                let rows = peers
+                    .into_iter()
+                    .map(|global_rank| {
+                        train_step_rows(
+                            16,
+                            source_shard_for_topology(&topology, global_rank, true).unwrap(),
+                        )
+                        .unwrap()
+                    })
+                    .collect::<Vec<_>>();
+                assert_eq!(rows[0], rows[1]);
+                assert_eq!(
+                    rows[0],
+                    (dp_rank * 8 + ep_rank * 4)..(dp_rank * 8 + ep_rank * 4 + 4)
+                );
+            }
         }
     }
+
+    #[test]
+    fn multi_lora_supports_replicated_and_global_dp_rows() {
+        let shard = SourceShard { rank: 1, size: 2 };
+        assert_eq!(multi_lora_rows(6, 3, Some(shard)).unwrap(), 3..6);
+        assert!(multi_lora_rows(3, 3, Some(shard)).is_err());
+        assert!(multi_lora_rows(1, 3, Some(shard)).is_err());
+        assert!(multi_lora_rows(9, 3, Some(shard)).is_err());
+        assert!(multi_lora_rows(6, 3, None).is_err());
+        assert_eq!(multi_lora_rows(3, 3, None).unwrap(), 0..3);
+        assert_eq!(multi_lora_rows(1, 3, None).unwrap(), 0..1);
+    }
+
+    #[test]
+    fn multi_lora_rows_support_explicit_pipeline_microbatches() {
+        assert_eq!(
+            multi_lora_rows_with_microbatches(
+                12,
+                3,
+                Some(SourceShard { rank: 1, size: 2 }),
+                2,
+            )
+            .unwrap(),
+            6..12
+        );
+        assert_eq!(
+            multi_lora_rows_with_microbatches(6, 3, None, 2).unwrap(),
+            0..6
+        );
+        assert!(multi_lora_rows_with_microbatches(9, 3, None, 2).is_err());
+        assert!(multi_lora_rows_with_microbatches(12, 3, Some(SourceShard { rank: 1, size: 2 }), 0)
+            .is_err());
+    }
+
+    #[test]
+    fn multi_lora_admission_uses_launcher_world_size() {
+        assert!(validate_multi_lora_global_batch_size(6, 3, 2).is_ok());
+        assert!(validate_multi_lora_global_batch_size(3, 3, 2).is_err());
+    }
+
+    #[test]
+    fn validates_flat_shapes_before_tensor_slicing() {
+        validate_flat_tensor_lengths(4, 8, 32, 32, 32).unwrap();
+        assert!(validate_flat_tensor_lengths(4, 8, 31, 32, 32).is_err());
+        assert!(validate_flat_tensor_lengths(0, 8, 0, 0, 0).is_err());
+        assert_eq!(tensor_element_range(4, 8, 1..3).unwrap(), 8..24);
+    }
+
+    #[test]
+    fn distributed_worker_group_enforces_singleton_session_ownership() {
+        let mut active = None;
+        assert!(require_worker_session(&active, "tenant-a").is_err());
+        create_worker_session(&mut active, "tenant-a").unwrap();
+        create_worker_session(&mut active, "tenant-a").unwrap();
+        require_worker_session(&active, "tenant-a").unwrap();
+        assert!(create_worker_session(&mut active, "tenant-b").is_err());
+        assert!(require_worker_session(&active, "tenant-b").is_err());
+        assert!(delete_worker_session(&mut active, "tenant-b").is_err());
+        delete_worker_session(&mut active, "tenant-a").unwrap();
+        create_worker_session(&mut active, "tenant-b").unwrap();
+        require_worker_session(&active, "tenant-b").unwrap();
+    }
+
+    #[test]
+    fn checkpoint_generation_is_unique_and_path_safe() {
+        let first = checkpoint_generation(42, 100, 0, "save").unwrap();
+        let second = checkpoint_generation(42, 100, 1, "save").unwrap();
+        assert_ne!(first, second);
+        assert!(first
+            .bytes()
+            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.')));
+        assert!(checkpoint_generation(42, 100, 2, "bad/path").is_err());
+        assert!(checkpoint_generation(42, 100, 3, "").is_err());
+    }
+
+    #[test]
+    fn checkpoint_metadata_compares_loss_bits_including_nan() {
+        let nan = f64::from_bits(0x7ff8_0000_0000_0042);
+        assert!(checkpoint_metadata_matches((7, nan), (7, nan)));
+        assert!(!checkpoint_metadata_matches((8, nan), (7, nan)));
+        assert!(!checkpoint_metadata_matches(
+            (7, nan),
+            (7, f64::from_bits(nan.to_bits() + 1))
+        ));
+    }
+
+    #[test]
+    fn checkpoint_staging_is_exclusive_and_sibling_scoped() {
+        let dir = tempfile::tempdir().unwrap();
+        let final_path = dir.path().join("checkpoint");
+        let generation = checkpoint_generation(42, 100, 0, "save").unwrap();
+        let expected = checkpoint_staging_path(&final_path, &generation).unwrap();
+        assert_eq!(expected.parent(), final_path.parent());
+
+        let staging = create_checkpoint_staging(&final_path, &generation).unwrap();
+        assert_eq!(staging, expected);
+        assert!(staging.is_dir());
+        assert!(create_checkpoint_staging(&final_path, &generation).is_err());
+    }
+
+    #[test]
+    fn checkpoint_publish_never_replaces_existing_destination() {
+        let dir = tempfile::tempdir().unwrap();
+        let final_path = dir.path().join("checkpoint");
+        let generation = checkpoint_generation(42, 100, 0, "save").unwrap();
+        let staging = create_checkpoint_staging(&final_path, &generation).unwrap();
+        std::fs::write(staging.join("receipt"), b"prepared").unwrap();
+        std::fs::create_dir(&final_path).unwrap();
+        std::fs::write(final_path.join("receipt"), b"existing").unwrap();
+
+        assert!(publish_checkpoint_noreplace(&staging, &final_path).is_err());
+        assert_eq!(
+            std::fs::read(final_path.join("receipt")).unwrap(),
+            b"existing"
+        );
+        assert!(staging.exists());
+
+        std::fs::remove_dir_all(&final_path).unwrap();
+        publish_checkpoint_noreplace(&staging, &final_path).unwrap();
+        assert!(!staging.exists());
+        assert_eq!(
+            std::fs::read(final_path.join("receipt")).unwrap(),
+            b"prepared"
+        );
+    }
 }
diff --git a/crates/rustrain-server/src/ep_dispatch.rs b/crates/rustrain-server/src/ep_dispatch.rs
new file mode 100644
index 00000000..f6515c41
--- /dev/null
+++ b/crates/rustrain-server/src/ep_dispatch.rs
@@ -0,0 +1,217 @@
+use std::fmt;
+
+use tokio::sync::{mpsc, oneshot};
+
+const DEFAULT_QUEUE_CAPACITY: usize = 32;
+const HARD_MAX_QUEUE_CAPACITY: usize = 4096;
+
+type DispatchJob = Box<dyn FnOnce() + Send + 'static>;
+
+#[derive(Clone)]
+pub(crate) struct EpDispatchScheduler {
+    sender: mpsc::Sender<DispatchJob>,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum EpDispatchScheduleError {
+    QueueFull,
+    QueueClosed,
+    WorkerFailed,
+}
+
+impl fmt::Display for EpDispatchScheduleError {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::QueueFull => formatter.write_str("EP dispatch queue is full"),
+            Self::QueueClosed => formatter.write_str("EP dispatch queue is closed"),
+            Self::WorkerFailed => formatter.write_str("EP dispatch worker failed"),
+        }
+    }
+}
+
+impl EpDispatchScheduler {
+    pub(crate) fn new(capacity: usize) -> Self {
+        assert!(capacity > 0, "EP dispatch queue capacity must be positive");
+        let (sender, mut receiver) = mpsc::channel::<DispatchJob>(capacity);
+        tokio::spawn(async move {
+            while let Some(job) = receiver.recv().await {
+                if let Err(error) = tokio::task::spawn_blocking(job).await {
+                    tracing::error!(%error, "EP dispatch job panicked");
+                }
+            }
+        });
+        Self { sender }
+    }
+
+    pub(crate) async fn run<T, F>(&self, operation: F) -> Result<T, EpDispatchScheduleError>
+    where
+        T: Send + 'static,
+        F: FnOnce() -> T + Send + 'static,
+    {
+        let receiver = self.submit(operation)?;
+        receiver
+            .await
+            .map_err(|_| EpDispatchScheduleError::WorkerFailed)
+    }
+
+    pub(crate) fn submit<T, F>(
+        &self,
+        operation: F,
+    ) -> Result<oneshot::Receiver<T>, EpDispatchScheduleError>
+    where
+        T: Send + 'static,
+        F: FnOnce() -> T + Send + 'static,
+    {
+        let (response, receiver) = oneshot::channel();
+        let job = Box::new(move || {
+            let _ = response.send(operation());
+        });
+        self.sender.try_send(job).map_err(|error| match error {
+            mpsc::error::TrySendError::Full(_) => EpDispatchScheduleError::QueueFull,
+            mpsc::error::TrySendError::Closed(_) => EpDispatchScheduleError::QueueClosed,
+        })?;
+        Ok(receiver)
+    }
+}
+
+pub(crate) fn configured_queue_capacity() -> usize {
+    std::env::var("RUSTRAIN_EP_DISPATCH_QUEUE_CAPACITY")
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .filter(|capacity| *capacity > 0)
+        .unwrap_or(DEFAULT_QUEUE_CAPACITY)
+        .min(HARD_MAX_QUEUE_CAPACITY)
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::{Arc, Condvar, Mutex};
+
+    use tokio::sync::oneshot;
+
+    use super::{EpDispatchScheduleError, EpDispatchScheduler};
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn accepted_jobs_are_fifo_and_single_inflight() {
+        let scheduler = EpDispatchScheduler::new(4);
+        let order = Arc::new(Mutex::new(Vec::new()));
+        let inflight = Arc::new(AtomicUsize::new(0));
+        let max_inflight = Arc::new(AtomicUsize::new(0));
+
+        let mut receivers = Vec::new();
+        for index in 0..3 {
+            let order = Arc::clone(&order);
+            let inflight = Arc::clone(&inflight);
+            let max_inflight = Arc::clone(&max_inflight);
+            receivers.push(
+                scheduler
+                    .submit(move || {
+                        let current = inflight.fetch_add(1, Ordering::SeqCst) + 1;
+                        max_inflight.fetch_max(current, Ordering::SeqCst);
+                        order.lock().unwrap().push(index);
+                        inflight.fetch_sub(1, Ordering::SeqCst);
+                        index
+                    })
+                    .unwrap(),
+            );
+        }
+
+        for (index, receiver) in receivers.into_iter().enumerate() {
+            assert_eq!(receiver.await.unwrap(), index);
+        }
+        assert_eq!(*order.lock().unwrap(), vec![0, 1, 2]);
+        assert_eq!(max_inflight.load(Ordering::SeqCst), 1);
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn full_queue_rejects_without_waiting() {
+        let scheduler = EpDispatchScheduler::new(1);
+        let gate = Arc::new((Mutex::new(false), Condvar::new()));
+        let (started, started_rx) = oneshot::channel();
+        let job_gate = Arc::clone(&gate);
+        let first = scheduler
+            .submit(move || {
+                let _ = started.send(());
+                let (lock, wake) = &*job_gate;
+                let mut released = lock.lock().unwrap();
+                while !*released {
+                    released = wake.wait(released).unwrap();
+                }
+            })
+            .unwrap();
+        started_rx.await.unwrap();
+
+        let second = scheduler.submit(|| 2usize).unwrap();
+        assert_eq!(
+            scheduler.submit(|| 3usize).unwrap_err(),
+            EpDispatchScheduleError::QueueFull
+        );
+
+        let (lock, wake) = &*gate;
+        *lock.lock().unwrap() = true;
+        wake.notify_one();
+        first.await.unwrap();
+        assert_eq!(second.await.unwrap(), 2);
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn panicking_job_does_not_stop_the_consumer() {
+        let scheduler = EpDispatchScheduler::new(2);
+        assert_eq!(
+            scheduler.run(|| panic!("injected dispatch panic")).await,
+            Err(EpDispatchScheduleError::WorkerFailed)
+        );
+        assert_eq!(scheduler.run(|| 7usize).await.unwrap(), 7);
+    }
+
+    #[tokio::test(flavor = "current_thread")]
+    async fn blocking_job_does_not_block_the_async_runtime() {
+        let scheduler = EpDispatchScheduler::new(1);
+        let gate = Arc::new((Mutex::new(false), Condvar::new()));
+        let job_gate = Arc::clone(&gate);
+        let (started, started_rx) = oneshot::channel();
+        let task = tokio::spawn(async move {
+            scheduler
+                .run(move || {
+                    let _ = started.send(());
+                    let (lock, wake) = &*job_gate;
+                    let mut released = lock.lock().unwrap();
+                    while !*released {
+                        released = wake.wait(released).unwrap();
+                    }
+                })
+                .await
+        });
+        started_rx.await.unwrap();
+
+        tokio::time::timeout(std::time::Duration::from_millis(100), async {
+            tokio::task::yield_now().await;
+        })
+        .await
+        .expect("blocking dispatch must not occupy the async runtime thread");
+
+        let (lock, wake) = &*gate;
+        *lock.lock().unwrap() = true;
+        wake.notify_one();
+        task.await.unwrap().unwrap();
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn dropping_response_does_not_cancel_an_accepted_job() {
+        let scheduler = EpDispatchScheduler::new(1);
+        let (completed, completed_rx) = oneshot::channel();
+        let response = scheduler
+            .submit(move || {
+                let _ = completed.send(());
+                9usize
+            })
+            .unwrap();
+        drop(response);
+
+        tokio::time::timeout(std::time::Duration::from_secs(1), completed_rx)
+            .await
+            .expect("accepted dispatch must execute after its caller disconnects")
+            .unwrap();
+    }
+}
diff --git a/crates/rustrain-server/src/grpc.rs b/crates/rustrain-server/src/grpc.rs
index 1d3f20bc..ffd90317 100644
--- a/crates/rustrain-server/src/grpc.rs
+++ b/crates/rustrain-server/src/grpc.rs
@@ -75,7 +75,7 @@ impl TrainService for TrainServiceImpl {
         let count = s
             .init_lora(InitLoRARequest {
                 rank: req.rank,
-                alpha: req.alpha as i64,
+                alpha: req.alpha,
                 target_layers: req.target_layers.iter().map(|&l| l as usize).collect(),
                 target_modules: req.target_modules,
                 lr: req.lr,
@@ -141,9 +141,7 @@ impl TrainService for TrainServiceImpl {
                 attention_mask,
             })
             .map_err(|e| Status::internal(e.to_string()))?;
-        Ok(Response::new(EvalStepResponse {
-            loss: result.loss,
-        }))
+        Ok(Response::new(EvalStepResponse { loss: result.loss }))
     }
 
     async fn save_checkpoint(
@@ -200,7 +198,7 @@ impl TrainService for TrainServiceImpl {
             .ok_or_else(|| Status::not_found("session not found"))?;
         let s = session.lock().await;
         let count = s
-            .export_adapter(&req.path)
+            .export_adapter(&req.path, Some(req.adapter_id))
             .map_err(|e| Status::internal(e.to_string()))?;
         Ok(Response::new(AdapterInfo {
             path: req.path,
@@ -208,7 +206,25 @@ impl TrainService for TrainServiceImpl {
         }))
     }
 
-    type StreamMetricsStream = std::pin::Pin<Box<dyn futures::Stream<Item = Result<StepMetric, Status>> + Send>>;
+    async fn import_adapter(
+        &self,
+        request: Request<ImportAdapterRequest>,
+    ) -> Result<Response<ImportAdapterResponse>, Status> {
+        let req = request.into_inner();
+        let session = self
+            .manager
+            .get_session(&req.session_id)
+            .await
+            .ok_or_else(|| Status::not_found("session not found"))?;
+        let mut s = session.lock().await;
+        let adapter_id = s
+            .import_adapter(&req.path)
+            .map_err(|e| Status::internal(e.to_string()))?;
+        Ok(Response::new(ImportAdapterResponse { adapter_id }))
+    }
+
+    type StreamMetricsStream =
+        std::pin::Pin<Box<dyn futures::Stream<Item = Result<StepMetric, Status>> + Send>>;
 
     async fn stream_metrics(
         &self,
@@ -274,6 +290,10 @@ impl TrainService for TrainServiceImpl {
                 alpha: req.alpha,
                 target_layers: req.target_layers,
                 target_modules: req.target_modules,
+                optimizer_lr: req.optimizer_lr.map(|value| value.value),
+                optimizer_beta1: req.optimizer_beta1.map(|value| value.value),
+                optimizer_beta2: req.optimizer_beta2.map(|value| value.value),
+                optimizer_eps: req.optimizer_eps.map(|value| value.value),
             })
             .map_err(|e| Status::internal(e.to_string()))?;
         Ok(Response::new(train::AddLoRaResponse { adapter_id: id }))
@@ -321,7 +341,12 @@ fn decode_tensor_data(td: &Option<TensorData>) -> Result<tch::Tensor, Status> {
         "int64" => tch::Kind::Int64,
         "float32" => tch::Kind::Float,
         "bfloat16" => tch::Kind::BFloat16,
-        _ => return Err(Status::invalid_argument(format!("unsupported dtype: {}", td.dtype))),
+        _ => {
+            return Err(Status::invalid_argument(format!(
+                "unsupported dtype: {}",
+                td.dtype
+            )));
+        }
     };
     let tensor = match kind {
         tch::Kind::Int64 => {
@@ -342,6 +367,11 @@ fn decode_tensor_data(td: &Option<TensorData>) -> Result<tch::Tensor, Status> {
         }
         _ => return Err(Status::invalid_argument("only int64 and float32 supported")),
     };
-    let local_rank = std::env::var("LOCAL_RANK").ok().and_then(|s| s.parse::<usize>().ok()).unwrap_or(0);
-    Ok(tensor.reshape(&td.shape).to_device(tch::Device::Cuda(local_rank)))
+    let local_rank = std::env::var("LOCAL_RANK")
+        .ok()
+        .and_then(|s| s.parse::<usize>().ok())
+        .unwrap_or(0);
+    Ok(tensor
+        .reshape(&td.shape)
+        .to_device(tch::Device::Cuda(local_rank)))
 }
diff --git a/crates/rustrain-server/src/lib.rs b/crates/rustrain-server/src/lib.rs
index 8e4a6556..09918b88 100644
--- a/crates/rustrain-server/src/lib.rs
+++ b/crates/rustrain-server/src/lib.rs
@@ -1,6 +1,7 @@
 pub mod api;
 pub mod checkpoint;
 pub mod ep;
+mod ep_dispatch;
 pub mod grpc;
 pub mod metrics;
 pub mod session;
diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs
index e115c794..363da3a7 100644
--- a/crates/rustrain-server/src/session.rs
+++ b/crates/rustrain-server/src/session.rs
@@ -1,6 +1,6 @@
 //! Training session trait + Qwen3.6 implementation.
 
-use anyhow::{anyhow, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use std::path::PathBuf;
 use std::sync::Arc;
 use tch::{Device, Kind, Tensor};
@@ -8,6 +8,296 @@ use tokio::sync::Mutex;
 
 use crate::checkpoint;
 use crate::metrics::{FileMetricsSink, MetricsSink, StepMetric};
+use rustrain_parallel::topology::ParallelTopology;
+use rustrain_qwen3_6::kernel::DynamicAdamConfig;
+use rustrain_qwen3_6::lora::{Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule};
+use rustrain_qwen3_6::pipeline::{stage_lora_slots, PipelineStageLayout};
+
+fn dynamic_pipeline_schedule(
+    pp_rank: usize,
+    pp_size: usize,
+    num_microbatches: usize,
+) -> Result<Vec<(Option<i64>, Option<i64>)>> {
+    if pp_size < 2 || pp_rank >= pp_size || num_microbatches == 0 {
+        bail!("invalid dynamic pipeline schedule PP={pp_rank}/{pp_size} microbatches={num_microbatches}");
+    }
+    let warmup = (pp_size - pp_rank - 1).min(num_microbatches);
+    let mut schedule = Vec::with_capacity(num_microbatches + warmup);
+    for mb in 0..warmup {
+        schedule.push((Some(mb as i64), None));
+    }
+    for mb in warmup..num_microbatches {
+        schedule.push((Some(mb as i64), Some((mb - warmup) as i64)));
+    }
+    for mb in num_microbatches - warmup..num_microbatches {
+        schedule.push((None, Some(mb as i64)));
+    }
+    Ok(schedule)
+}
+
+pub(crate) fn configured_dynamic_pipeline_microbatches() -> Result<usize> {
+    let raw = std::env::var("QWEN36_PP_MICROBATCHES").unwrap_or_else(|_| "1".to_string());
+    let count = raw
+        .parse::<usize>()
+        .with_context(|| format!("QWEN36_PP_MICROBATCHES must be a positive integer, got {raw}"))?;
+    if count == 0 || count > 4096 {
+        bail!("QWEN36_PP_MICROBATCHES must be in 1..=4096, got {count}");
+    }
+    Ok(count)
+}
+
+fn prepare_dynamic_pipeline_microbatches(
+    input: &TrainInput,
+    adapter_count: usize,
+) -> Result<Vec<(Tensor, Tensor, Tensor)>> {
+    if adapter_count == 0 {
+        bail!("dynamic pipeline requires at least one selected adapter");
+    }
+    let input_shape = input.input_ids.size();
+    if input_shape.len() != 2 {
+        bail!("dynamic pipeline input_ids must be rank-2");
+    }
+    if input.target_mask.size() != input_shape || input.attention_mask.size() != input_shape {
+        bail!("dynamic pipeline input and mask shapes must match");
+    }
+    let input_batch = input_shape[0];
+    if input_batch <= 0 {
+        bail!("dynamic pipeline input batch must be positive");
+    }
+    let configured = configured_dynamic_pipeline_microbatches()?;
+    let microbatch_count = if input_batch == 1 { 1 } else { configured };
+    let expected_rows = i64::try_from(
+        adapter_count
+            .checked_mul(microbatch_count)
+            .ok_or_else(|| anyhow!("dynamic pipeline microbatch row count overflowed"))?,
+    )
+    .context("dynamic pipeline microbatch row count exceeds i64")?;
+    if input_batch != 1 && input_batch != expected_rows {
+        bail!(
+            "dynamic pipeline input batch={input_batch} must be 1 or selected_adapters*microbatches={expected_rows}"
+        );
+    }
+    if microbatch_count == 1 && input_batch == 1 {
+        return Ok(vec![(
+            input.input_ids.shallow_clone(),
+            input.target_mask.shallow_clone(),
+            input.attention_mask.shallow_clone(),
+        )]);
+    }
+
+    let rows_per_microbatch = if input_batch == 1 {
+        1
+    } else {
+        i64::try_from(adapter_count).context("adapter count exceeds i64")?
+    };
+    let mut microbatches = Vec::with_capacity(microbatch_count);
+    for index in 0..microbatch_count {
+        let start = i64::try_from(index)
+            .context("dynamic pipeline microbatch index exceeds i64")?
+            * rows_per_microbatch;
+        microbatches.push((
+            input.input_ids.narrow(0, start, rows_per_microbatch),
+            input.target_mask.narrow(0, start, rows_per_microbatch),
+            input.attention_mask.narrow(0, start, rows_per_microbatch),
+        ));
+    }
+    Ok(microbatches)
+}
+
+fn validate_qwen_parallel_features(
+    is_moe: bool,
+    tp_size: usize,
+    ep_size: usize,
+    ep_a2a: bool,
+    ep_a2a_sharded: bool,
+) -> Result<()> {
+    let is_ep = is_moe && ep_size > 1;
+    if ep_a2a_sharded && !is_ep {
+        bail!("QWEN36_EP_A2A_SHARDED=1 requires expert-parallel training");
+    }
+    if ep_a2a_sharded && !ep_a2a {
+        bail!("QWEN36_EP_A2A_SHARDED=1 requires QWEN36_EP_A2A=1");
+    }
+    if is_moe && tp_size > 1 && ep_size > 1 && !ep_a2a_sharded {
+        bail!(
+            "Qwen server MoE TPxEP requires QWEN36_EP_A2A_SHARDED=1; replicated expert TP is not supported by the native kernel"
+        );
+    }
+    Ok(())
+}
+
+fn validate_tp_intermediate_sizes(tp_size: usize, sizes: &[(&str, i64)]) -> Result<()> {
+    for (kind, intermediate) in sizes {
+        if *intermediate <= 0 || *intermediate % tp_size as i64 != 0 {
+            bail!("{kind} intermediate_size={intermediate} must be divisible by TP_SIZE={tp_size}");
+        }
+    }
+    Ok(())
+}
+
+struct ResolvedQwenTopology {
+    topology: ParallelTopology,
+    global_rank: usize,
+    world_size: usize,
+    local_rank: usize,
+    tp_rank: usize,
+    cp_rank: usize,
+    ep_rank: usize,
+    dp_rank: usize,
+    pp_rank: usize,
+    tp_size: usize,
+    cp_size: usize,
+    ep_size: usize,
+    dp_size: usize,
+    pp_size: usize,
+}
+
+fn resolve_qwen_topology(
+    runtime_config: &rustrain_qwen3_6::config::Qwen36RuntimeConfig,
+) -> Result<ResolvedQwenTopology> {
+    let global_rank = std::env::var("RANK")
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(0);
+    let world_size = std::env::var("WORLD_SIZE")
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(1);
+    let local_rank = std::env::var("LOCAL_RANK")
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(0);
+    let tp_size = std::env::var("TP_SIZE")
+        .or_else(|_| std::env::var("RUSTRAIN_TP_SIZE"))
+        .or_else(|_| std::env::var("TENSOR_MODEL_PARALLEL_SIZE"))
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(1);
+    let has_explicit_non_tp_axis = [
+        "EP_SIZE",
+        "RUSTRAIN_EP_SIZE",
+        "EXPERT_MODEL_PARALLEL_SIZE",
+        "DP_SIZE",
+        "RUSTRAIN_DP_SIZE",
+        "DATA_PARALLEL_SIZE",
+        "PP_SIZE",
+        "RUSTRAIN_PP_SIZE",
+        "PIPELINE_MODEL_PARALLEL_SIZE",
+        "CP_SIZE",
+        "RUSTRAIN_CP_SIZE",
+        "CONTEXT_PARALLEL_SIZE",
+    ]
+    .iter()
+    .any(|name| std::env::var_os(name).is_some());
+    let topology = if runtime_config.is_moe && !has_explicit_non_tp_axis {
+        let ep_size = world_size
+            .checked_div(tp_size)
+            .filter(|size| size * tp_size == world_size)
+            .ok_or_else(|| {
+                anyhow!("WORLD_SIZE={world_size} is not divisible by TP_SIZE={tp_size}")
+            })?;
+        let rank_order = std::env::var("RUSTRAIN_PARALLEL_ORDER")
+            .or_else(|_| std::env::var("PARALLEL_ORDER"))
+            .unwrap_or_else(|_| rustrain_parallel::topology::DEFAULT_RANK_ORDER.to_string());
+        ParallelTopology::with_order(tp_size, 1, 1, ep_size, 1, &rank_order)?
+    } else {
+        ParallelTopology::from_env_with_world_size(world_size)?
+    };
+    let cp_full_attention = std::env::var("QWEN36_CP_FULL_ATTENTION_KV_GATHER")
+        .map(|value| !value.is_empty() && value != "0")
+        .unwrap_or(false);
+    if topology.context_parallel_size() != 1
+        && (topology.context_parallel_size() != 2
+            || topology.tensor_model_parallel_size() != 1
+            || topology.pipeline_model_parallel_size() != 1
+            || topology.data_parallel_size() != 1
+            || topology.expert_model_parallel_size() != 1
+            || !cp_full_attention)
+    {
+        bail!(
+            "native Qwen server CP requires CP2 with TP=EP=DP=PP=1 and QWEN36_CP_FULL_ATTENTION_KV_GATHER=1 (tp={} pp={} dp={} ep={} cp={})",
+            topology.tensor_model_parallel_size(),
+            topology.pipeline_model_parallel_size(),
+            topology.data_parallel_size(),
+            topology.expert_model_parallel_size(),
+            topology.context_parallel_size(),
+        );
+    }
+    if topology.tensor_model_parallel_size() != tp_size {
+        bail!(
+            "Qwen server topology TP={} does not match TP_SIZE={tp_size}",
+            topology.tensor_model_parallel_size()
+        );
+    }
+    topology.coordinates(global_rank)?;
+    if !runtime_config.is_moe && topology.expert_model_parallel_size() != 1 {
+        bail!("native dense Qwen server requires expert_model_parallel_size=1");
+    }
+
+    Ok(ResolvedQwenTopology {
+        tp_rank: topology.tensor_rank(global_rank)?,
+        cp_rank: topology.context_rank(global_rank)?,
+        ep_rank: topology.expert_rank(global_rank)?,
+        dp_rank: topology.data_rank(global_rank)?,
+        pp_rank: topology.pipeline_rank(global_rank)?,
+        tp_size: topology.tensor_model_parallel_size(),
+        cp_size: topology.context_parallel_size(),
+        ep_size: topology.expert_model_parallel_size(),
+        dp_size: topology.data_parallel_size(),
+        pp_size: topology.pipeline_model_parallel_size(),
+        topology,
+        global_rank,
+        world_size,
+        local_rank,
+    })
+}
+
+fn validate_dynamic_adapter_manifests<'a>(
+    manifests: impl IntoIterator<Item = &'a checkpoint::DynamicAdapterManifest>,
+) -> Result<()> {
+    let mut adapter_ids = std::collections::BTreeSet::new();
+    for manifest in manifests {
+        if manifest.id <= 0 || !adapter_ids.insert(manifest.id) {
+            bail!("checkpoint dynamic adapter IDs must be positive and unique");
+        }
+        if i64::try_from(manifest.optimizer_step).is_err() {
+            bail!(
+                "dynamic adapter {} optimizer step exceeds native range",
+                manifest.id
+            );
+        }
+        if manifest
+            .optimizer_lr
+            .is_some_and(|optimizer_lr| !optimizer_lr.is_finite() || optimizer_lr < 0.0)
+        {
+            bail!(
+                "dynamic adapter {} optimizer learning rate must be finite and non-negative",
+                manifest.id
+            );
+        }
+        for (name, value) in [
+            ("beta1", manifest.optimizer_beta1),
+            ("beta2", manifest.optimizer_beta2),
+        ] {
+            if value.is_some_and(|value| !value.is_finite() || !(0.0..1.0).contains(&value)) {
+                bail!(
+                    "dynamic adapter {} optimizer {name} must be finite and in [0, 1)",
+                    manifest.id
+                );
+            }
+        }
+        if manifest
+            .optimizer_eps
+            .is_some_and(|value| !value.is_finite() || value < 0.0)
+        {
+            bail!(
+                "dynamic adapter {} optimizer epsilon must be finite and non-negative",
+                manifest.id
+            );
+        }
+    }
+    Ok(())
+}
 
 /// Session states.
 #[derive(Debug, Clone)]
@@ -33,10 +323,10 @@ pub struct SessLoadDatasetRequest {
     pub seq_len: usize,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct InitLoRARequest {
     pub rank: i64,
-    pub alpha: i64,
+    pub alpha: f64,
     pub target_layers: Vec<usize>,
     pub target_modules: Vec<String>,
     pub lr: f64,
@@ -58,11 +348,24 @@ pub struct TrainOutput {
     pub step: u64,
 }
 
+#[derive(Debug)]
+pub struct MultiLoraTrainOutput {
+    pub loss: f64,
+    pub adapter_losses: Vec<f64>,
+    pub adapter_steps: Vec<u64>,
+    pub step: u64,
+}
+
 #[derive(Debug)]
 pub struct EvalOutput {
     pub loss: f64,
 }
 
+#[derive(Debug)]
+pub struct MultiLoraEvalOutput {
+    pub adapter_losses: Vec<(i64, f64)>,
+}
+
 #[derive(Debug)]
 pub struct SessionStatus {
     pub state: String,
@@ -77,11 +380,59 @@ pub trait TrainingSession: Send {
     fn load_dataset(&mut self, req: SessLoadDatasetRequest) -> Result<usize>;
     fn init_lora(&mut self, req: InitLoRARequest) -> Result<usize>;
     fn train_step(&mut self, input: TrainInput) -> Result<TrainOutput>;
-    fn train_multi_lora(&mut self, input: TrainInput, n_total: i32, rank: i32) -> Result<TrainOutput>;
+    fn train_multi_lora(
+        &mut self,
+        input: TrainInput,
+        n_total: i32,
+        rank: i32,
+        adapter_ids: &[i64],
+    ) -> Result<TrainOutput>;
     fn eval_step(&self, input: TrainInput) -> Result<EvalOutput>;
+    fn validate_dynamic_adapter_steps(
+        &self,
+        adapter_ids: &[i64],
+        expected_steps: &[u64],
+    ) -> Result<()> {
+        let _ = (adapter_ids, expected_steps);
+        bail!("dynamic adapter step validation is not supported by this session")
+    }
+    fn eval_multi_lora_host_i64(
+        &self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+        adapter_ids: &[i64],
+    ) -> Result<MultiLoraEvalOutput> {
+        let _ = (
+            input_ids,
+            target_mask,
+            attention_mask,
+            batch_size,
+            seq_len,
+            adapter_ids,
+        );
+        bail!("selected multi-LoRA evaluation is not supported by this session")
+    }
     fn save_checkpoint(&self, path: &str) -> Result<(u64, f64)>;
+    fn save_checkpoint_with_generation(
+        &self,
+        path: &str,
+        checkpoint_generation: Option<&str>,
+    ) -> Result<(u64, f64)> {
+        if checkpoint_generation.is_some() {
+            bail!("this training session does not support coordinated checkpoint generations");
+        }
+        self.save_checkpoint(path)
+    }
     fn load_checkpoint(&mut self, path: &str) -> Result<(u64, f64)>;
-    fn export_adapter(&self, path: &str) -> Result<usize>;
+    #[doc(hidden)]
+    fn load_checkpoint_in_place(&mut self, path: &str) -> Result<(u64, f64)>;
+    fn export_adapter(&self, path: &str, adapter_id: Option<i64>) -> Result<usize>;
+    /// Import one PEFT-style adapter as a new dynamic adapter. The fixed
+    /// adapter remains reserved as ID 0 and is never overwritten by import.
+    fn import_adapter(&mut self, path: &str) -> Result<i64>;
     fn status(&self) -> SessionStatus;
     fn get_metrics(&self) -> Vec<StepMetric>;
 
@@ -96,12 +447,62 @@ pub trait TrainingSession: Send {
     fn list_lora(&self) -> Vec<i64>;
 }
 
+#[derive(Clone)]
+struct QwenContextSpec {
+    runtime_config: rustrain_qwen3_6::config::Qwen36RuntimeConfig,
+    stage: PipelineStageLayout,
+    init: InitLoRARequest,
+    target_layers: Vec<usize>,
+    target_modules: Vec<Qwen36LoraTargetModule>,
+    base_tp_attention: bool,
+    base_tp_mlp: bool,
+    vocab_parallel: bool,
+    data_parallel: bool,
+    expert_parallel: bool,
+    expert_start: usize,
+    expert_count: usize,
+    global_rank: usize,
+    world_size: usize,
+    tp_rank: usize,
+    tp_size: usize,
+    tp_color: usize,
+    cp_rank: usize,
+    cp_size: usize,
+    cp_color: usize,
+    ep_rank: usize,
+    ep_size: usize,
+    ep_color: usize,
+    dp_rank: usize,
+    dp_size: usize,
+    dp_color: usize,
+    pp_rank: usize,
+    pp_size: usize,
+    pp_color: usize,
+}
+
+struct PendingCheckpointLoad {
+    transaction_id: String,
+    source_path: String,
+    ctx: rustrain_qwen3_6::kernel::CppTrainingContext,
+    dynamic_lora_configs: std::collections::BTreeMap<i64, Qwen36LoraConfig>,
+    dynamic_lora_optimizer_configs: std::collections::BTreeMap<i64, DynamicAdamConfig>,
+    state: SessionState,
+    step: u64,
+    last_loss: f64,
+}
+
 #[derive(Debug)]
 pub struct AddLoRARequest {
     pub rank: i64,
     pub alpha: f64,
     pub target_layers: Vec<i64>,
-    pub target_modules: String,  // comma-separated, empty = all
+    pub target_modules: String, // comma-separated, empty = all
+    /// `None` inherits the training context learning rate.
+    pub optimizer_lr: Option<f64>,
+    /// Missing Adam scalars inherit the training context defaults.
+    pub optimizer_beta1: Option<f64>,
+    pub optimizer_beta2: Option<f64>,
+    pub optimizer_eps: Option<f64>,
 }
 
 /// Qwen3.6 training session — wraps CppTrainingContext.
@@ -116,11 +517,20 @@ pub struct Qwen36Session {
     weights: Option<std::collections::BTreeMap<String, Tensor>>,
     dataset: Option<rustrain_qwen3_6::sft::SftDataset>,
     lora_rank: i64,
-    lora_alpha: i64,
+    lora_alpha: f64,
+    lora_target_layers: Vec<usize>,
+    lora_target_modules: Vec<Qwen36LoraTargetModule>,
+    dynamic_lora_configs: std::collections::BTreeMap<i64, Qwen36LoraConfig>,
+    dynamic_lora_optimizer_configs: std::collections::BTreeMap<i64, DynamicAdamConfig>,
     lr: f64,
+    beta1: f64,
+    beta2: f64,
+    eps: f64,
     metrics: Option<Arc<FileMetricsSink>>,
     last_loss: f64,
     step: u64,
+    context_spec: Option<QwenContextSpec>,
+    pending_checkpoint_load: Option<PendingCheckpointLoad>,
     // NCCL marker for EP (comm stored in C++ TrainingContext)
     _nccl_ep: bool,
 }
@@ -130,6 +540,90 @@ pub struct Qwen36Session {
 // The Mutex in SessionManager ensures single-threaded access.
 unsafe impl Send for Qwen36Session {}
 
+impl Drop for Qwen36Session {
+    fn drop(&mut self) {
+        // Both native contexts borrow the frozen tensors through raw pointers.
+        // Destroy every context explicitly before Rust releases those tensors.
+        self.pending_checkpoint_load = None;
+        self.ctx = None;
+        self.weights = None;
+    }
+}
+
+fn create_qwen_context(
+    weights: &std::collections::BTreeMap<String, Tensor>,
+    compute_kind: Kind,
+    spec: &QwenContextSpec,
+    synchronize_parameters: bool,
+) -> Result<rustrain_qwen3_6::kernel::CppTrainingContext> {
+    let lora_scaling = spec.init.alpha / spec.init.rank as f64;
+    let ctx = rustrain_qwen3_6::kernel::CppTrainingContext::new_for_stage(
+        weights,
+        &spec.runtime_config,
+        &spec.stage,
+        compute_kind,
+        spec.init.lr,
+        spec.init.beta1,
+        spec.init.beta2,
+        spec.init.eps,
+        lora_scaling,
+        spec.init.rank,
+        spec.base_tp_attention,
+        spec.base_tp_mlp,
+        spec.vocab_parallel,
+        spec.data_parallel,
+        spec.expert_parallel,
+        &spec.target_layers,
+        &spec.target_modules,
+        spec.expert_start,
+        spec.expert_count,
+    )?;
+    if spec.world_size > 1 {
+        if synchronize_parameters {
+            ctx.init_parallel_nccl(
+                spec.global_rank,
+                spec.world_size,
+                spec.tp_rank,
+                spec.tp_size,
+                spec.tp_color,
+                spec.cp_rank,
+                spec.cp_size,
+                spec.cp_color,
+                spec.ep_rank,
+                spec.ep_size,
+                spec.ep_color,
+                spec.dp_rank,
+                spec.dp_size,
+                spec.dp_color,
+                spec.pp_rank,
+                spec.pp_size,
+                spec.pp_color,
+            )?;
+        } else {
+            ctx.attach_parallel_nccl_no_sync(
+                spec.global_rank,
+                spec.world_size,
+                spec.tp_rank,
+                spec.tp_size,
+                spec.tp_color,
+                spec.cp_rank,
+                spec.cp_size,
+                spec.cp_color,
+                spec.ep_rank,
+                spec.ep_size,
+                spec.ep_color,
+                spec.dp_rank,
+                spec.dp_size,
+                spec.dp_color,
+                spec.pp_rank,
+                spec.pp_size,
+                spec.pp_color,
+            )?;
+        }
+    }
+    Ok(ctx)
+}
+
 impl Qwen36Session {
     pub fn new(device: Device, compute_kind: Kind, metrics_path: PathBuf) -> Self {
         Self {
@@ -142,19 +636,571 @@ impl Qwen36Session {
             weights: None,
             dataset: None,
             lora_rank: 0,
-            lora_alpha: 0,
+            lora_alpha: 0.0,
+            lora_target_layers: Vec::new(),
+            lora_target_modules: Vec::new(),
+            dynamic_lora_configs: std::collections::BTreeMap::new(),
+            dynamic_lora_optimizer_configs: std::collections::BTreeMap::new(),
             lr: 1e-4,
+            beta1: 0.9,
+            beta2: 0.999,
+            eps: 1e-8,
             metrics: Some(Arc::new(FileMetricsSink::new(metrics_path))),
             last_loss: 0.0,
             step: 0,
+            context_spec: None,
+            pending_checkpoint_load: None,
             _nccl_ep: false,
         }
     }
 
+    fn default_dynamic_adam_config(&self) -> DynamicAdamConfig {
+        DynamicAdamConfig {
+            lr: self.lr,
+            beta1: self.beta1,
+            beta2: self.beta2,
+            eps: self.eps,
+        }
+    }
+
     /// Get the device this session is bound to.
     pub fn device(&self) -> Device {
         self.device
     }
+
+    pub fn native_context_is_healthy(&self) -> bool {
+        self.ctx
+            .as_ref()
+            .map(|ctx| ctx.is_healthy())
+            .unwrap_or(true)
+    }
+
+    fn train_multi_lora_pipeline(
+        &mut self,
+        input: &TrainInput,
+        n_total: i32,
+        adapter_ids: &[i64],
+    ) -> Result<TrainOutput> {
+        let (pp_rank, pp_size) = self
+            .context_spec
+            .as_ref()
+            .map(|spec| (spec.pp_rank, spec.pp_size))
+            .ok_or_else(|| anyhow!("Qwen context is not initialized"))?;
+        if pp_size < 2 {
+            bail!("dynamic pipeline training requires PP_SIZE >= 2");
+        }
+        let ctx = self
+            .ctx
+            .as_ref()
+            .ok_or_else(|| anyhow!("LoRA not initialized"))?;
+        let selected = if adapter_ids.is_empty() {
+            ctx.list_lora()
+                .into_iter()
+                .filter(|id| *id > 0)
+                .collect::<Vec<_>>()
+        } else {
+            adapter_ids.to_vec()
+        };
+        if selected.is_empty() || selected.len() != n_total as usize {
+            bail!("dynamic pipeline adapter count does not match n_total={n_total}");
+        }
+        let microbatches = prepare_dynamic_pipeline_microbatches(input, selected.len())?;
+        let num_microbatches = i64::try_from(microbatches.len())
+            .context("dynamic pipeline microbatch count exceeds i64")?;
+        let window_id = i64::try_from(self.step).context("pipeline step exceeds i64")?;
+        ctx.pipeline_begin_dynamic_selected_v1(window_id, num_microbatches, &selected)?;
+        let result = (|| -> Result<f64> {
+            for (forward_mb, backward_mb) in
+                dynamic_pipeline_schedule(pp_rank, pp_size, microbatches.len())?
+            {
+                let forward_input_ids = forward_mb
+                    .and_then(|index| microbatches.get(index as usize).map(|batch| &batch.0));
+                let forward_target_mask = forward_mb
+                    .and_then(|index| microbatches.get(index as usize).map(|batch| &batch.1));
+                let forward_attention_mask = forward_mb
+                    .and_then(|index| microbatches.get(index as usize).map(|batch| &batch.2));
+                let tick_result = ctx.pipeline_tick_v1(
+                    window_id,
+                    forward_mb,
+                    backward_mb,
+                    forward_input_ids,
+                    forward_target_mask,
+                    forward_attention_mask,
+                    1.0,
+                )?;
+                if !tick_result.loss.is_finite() {
+                    bail!("dynamic pipeline returned a non-finite tick loss");
+                }
+            }
+            Ok(ctx.pipeline_finish_v1()?.loss)
+        })();
+        match result {
+            Ok(loss) => self.finish_train_step(loss, false),
+            Err(error) => {
+                let _ = ctx.pipeline_abort_v1();
+                Err(error)
+            }
+        }
+    }
+
+    fn train_multi_lora_pipeline_report(
+        &mut self,
+        input: &TrainInput,
+        n_total: i32,
+        adapter_ids: &[i64],
+    ) -> Result<MultiLoraTrainOutput> {
+        let (pp_rank, pp_size) = self
+            .context_spec
+            .as_ref()
+            .map(|spec| (spec.pp_rank, spec.pp_size))
+            .ok_or_else(|| anyhow!("Qwen context is not initialized"))?;
+        if pp_size < 2 {
+            bail!("dynamic pipeline loss reporting requires PP_SIZE >= 2");
+        }
+        let ctx = self
+            .ctx
+            .as_ref()
+            .ok_or_else(|| anyhow!("LoRA not initialized"))?;
+        if adapter_ids.is_empty() || adapter_ids.len() != n_total as usize {
+            bail!("dynamic pipeline report adapter count does not match n_total={n_total}");
+        }
+        let microbatches = prepare_dynamic_pipeline_microbatches(input, adapter_ids.len())?;
+        let num_microbatches = i64::try_from(microbatches.len())
+            .context("dynamic pipeline microbatch count exceeds i64")?;
+        let window_id = i64::try_from(self.step).context("pipeline step exceeds i64")?;
+        ctx.pipeline_begin_dynamic_selected_v1(window_id, num_microbatches, adapter_ids)?;
+        let result = (|| -> Result<(f64, Vec<f64>)> {
+            for (forward_mb, backward_mb) in
+                dynamic_pipeline_schedule(pp_rank, pp_size, microbatches.len())?
+            {
+                let forward_input_ids = forward_mb
+                    .and_then(|index| microbatches.get(index as usize).map(|batch| &batch.0));
+                let forward_target_mask = forward_mb
+                    .and_then(|index| microbatches.get(index as usize).map(|batch| &batch.1));
+                let forward_attention_mask = forward_mb
+                    .and_then(|index| microbatches.get(index as usize).map(|batch| &batch.2));
+                let tick_result = ctx.pipeline_tick_v1(
+                    window_id,
+                    forward_mb,
+                    backward_mb,
+                    forward_input_ids,
+                    forward_target_mask,
+                    forward_attention_mask,
+                    1.0,
+                )?;
+                if !tick_result.loss.is_finite() {
+                    bail!("dynamic pipeline returned a non-finite tick loss");
+                }
+            }
+            let (finish, adapter_losses) =
+                ctx.pipeline_finish_dynamic_report_v1(adapter_ids.len())?;
+            Ok((finish.loss, adapter_losses))
+        })();
+        match result {
+            Ok((loss, adapter_losses)) => {
+                let adapter_steps = adapter_ids
+                    .iter()
+                    .map(|adapter_id| {
+                        u64::try_from(ctx.get_adapter_step_count(*adapter_id)?)
+                            .context("native dynamic adapter optimizer step is negative")
+                    })
+                    .collect::<Result<Vec<_>>>()?;
+                let output = self.finish_train_step(loss, false)?;
+                Ok(MultiLoraTrainOutput {
+                    loss: output.loss,
+                    adapter_losses,
+                    adapter_steps,
+                    step: output.step,
+                })
+            }
+            Err(error) => {
+                let _ = ctx.pipeline_abort_v1();
+                Err(error)
+            }
+        }
+    }
+
+    pub fn train_step_host_i64(
+        &mut self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+    ) -> Result<TrainOutput> {
+        let loss = self
+            .ctx
+            .as_ref()
+            .ok_or_else(|| anyhow!("LoRA not initialized"))?
+            .train_step_host_i64(input_ids, target_mask, attention_mask, batch_size, seq_len)?;
+        self.finish_train_step(loss, true)
+    }
+
+    pub fn train_multi_lora_host_i64(
+        &mut self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+        n_total: i32,
+        lora_rank: i32,
+        adapter_ids: &[i64],
+    ) -> Result<TrainOutput> {
+        if self
+            .context_spec
+            .as_ref()
+            .map(|spec| spec.pp_size > 1)
+            .unwrap_or(false)
+        {
+            let input = TrainInput {
+                input_ids: Tensor::from_slice(input_ids)
+                    .reshape([batch_size as i64, seq_len as i64])
+                    .to_device(self.device),
+                target_mask: Tensor::from_slice(target_mask)
+                    .reshape([batch_size as i64, seq_len as i64])
+                    .to_device(self.device),
+                attention_mask: Tensor::from_slice(attention_mask)
+                    .reshape([batch_size as i64, seq_len as i64])
+                    .to_device(self.device),
+            };
+            return self.train_multi_lora_pipeline(&input, n_total, adapter_ids);
+        }
+        let loss = self
+            .ctx
+            .as_ref()
+            .ok_or_else(|| anyhow!("LoRA not initialized"))?
+            .train_multi_lora_host_i64(
+                input_ids,
+                target_mask,
+                attention_mask,
+                batch_size,
+                seq_len,
+                n_total,
+                lora_rank,
+                adapter_ids,
+            )?;
+        self.finish_train_step(loss, false)
+    }
+
+    pub fn train_multi_lora_host_i64_report(
+        &mut self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+        n_total: i32,
+        lora_rank: i32,
+        adapter_ids: &[i64],
+    ) -> Result<MultiLoraTrainOutput> {
+        if self
+            .context_spec
+            .as_ref()
+            .map(|spec| spec.pp_size > 1)
+            .unwrap_or(false)
+        {
+            let input = TrainInput {
+                input_ids: Tensor::from_slice(input_ids)
+                    .reshape([batch_size as i64, seq_len as i64])
+                    .to_device(self.device),
+                target_mask: Tensor::from_slice(target_mask)
+                    .reshape([batch_size as i64, seq_len as i64])
+                    .to_device(self.device),
+                attention_mask: Tensor::from_slice(attention_mask)
+                    .reshape([batch_size as i64, seq_len as i64])
+                    .to_device(self.device),
+            };
+            return self.train_multi_lora_pipeline_report(&input, n_total, adapter_ids);
+        }
+        let ctx = self
+            .ctx
+            .as_ref()
+            .ok_or_else(|| anyhow!("LoRA not initialized"))?;
+        let report = ctx.train_multi_lora_host_i64_report(
+            input_ids,
+            target_mask,
+            attention_mask,
+            batch_size,
+            seq_len,
+            n_total,
+            lora_rank,
+            adapter_ids,
+        )?;
+        let adapter_steps = adapter_ids
+            .iter()
+            .map(|adapter_id| {
+                u64::try_from(ctx.get_adapter_step_count(*adapter_id)?)
+                    .context("native dynamic adapter optimizer step is negative")
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let loss = report.aggregate_loss;
+        let output = self.finish_train_step(loss, false)?;
+        Ok(MultiLoraTrainOutput {
+            loss: output.loss,
+            adapter_losses: report.adapter_losses,
+            adapter_steps,
+            step: output.step,
+        })
+    }
+
+    pub fn eval_step_host_i64(
+        &self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+    ) -> Result<EvalOutput> {
+        let loss = self
+            .ctx
+            .as_ref()
+            .ok_or_else(|| anyhow!("LoRA not initialized"))?
+            .eval_step_host_i64(input_ids, target_mask, attention_mask, batch_size, seq_len)?;
+        Ok(EvalOutput { loss })
+    }
+
+    pub fn eval_multi_lora_host_i64(
+        &self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+        adapter_ids: &[i64],
+    ) -> Result<MultiLoraEvalOutput> {
+        let losses = self
+            .ctx
+            .as_ref()
+            .ok_or_else(|| anyhow!("LoRA not initialized"))?
+            .eval_multi_lora_host_i64(
+                input_ids,
+                target_mask,
+                attention_mask,
+                batch_size,
+                seq_len,
+                adapter_ids,
+            )?;
+        Ok(MultiLoraEvalOutput {
+            adapter_losses: adapter_ids.iter().copied().zip(losses).collect(),
+        })
+    }
+
+    pub fn validate_dynamic_adapter_steps(
+        &self,
+        adapter_ids: &[i64],
+        expected_steps: &[u64],
+    ) -> Result<()> {
+        if expected_steps.is_empty() {
+            return Ok(());
+        }
+        if adapter_ids.is_empty() || expected_steps.len() != adapter_ids.len() {
+            bail!(
+                "expected_steps length {} must match adapter_ids length {}",
+                expected_steps.len(),
+                adapter_ids.len()
+            );
+        }
+        self.ctx
+            .as_ref()
+            .ok_or_else(|| anyhow!("LoRA not initialized"))?
+            .validate_adapter_steps(adapter_ids, expected_steps)
+    }
+
+    fn finish_train_step(&mut self, loss: f64, record_metric: bool) -> Result<TrainOutput> {
+        self.step = self
+            .step
+            .checked_add(1)
+            .context("training step counter overflow")?;
+        self.last_loss = loss;
+        self.state = SessionState::Training { step: self.step };
+        if record_metric {
+            if let Some(ref metrics) = self.metrics {
+                let mem_gb = rustrain_train::metrics::gpu_memory_allocated_mb()
+                    .map(|m| m / 1024.0)
+                    .unwrap_or(0.0);
+                metrics.record_step(StepMetric {
+                    step: self.step,
+                    loss,
+                    lr: self.lr,
+                    mem_gb,
+                    timestamp_unix: chrono::Utc::now().timestamp(),
+                });
+            }
+        }
+        Ok(TrainOutput {
+            loss,
+            step: self.step,
+        })
+    }
+
+    pub fn prepare_checkpoint_load(
+        &mut self,
+        path: &str,
+        transaction_id: &str,
+    ) -> Result<(u64, f64)> {
+        if transaction_id.is_empty()
+            || !transaction_id
+                .bytes()
+                .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
+        {
+            bail!("checkpoint transaction ID must be non-empty and path-safe");
+        }
+        if let Some(pending) = &self.pending_checkpoint_load {
+            if pending.transaction_id == transaction_id {
+                if pending.source_path != path {
+                    bail!(
+                        "checkpoint transaction {transaction_id} is already prepared from {}",
+                        pending.source_path
+                    );
+                }
+                return Ok((pending.step, pending.last_loss));
+            }
+            bail!(
+                "checkpoint transaction {} is already prepared",
+                pending.transaction_id
+            );
+        }
+        let spec = self
+            .context_spec
+            .as_ref()
+            .context("LoRA context specification is unavailable")?
+            .clone();
+        let weights = self
+            .weights
+            .as_ref()
+            .context("base model weights are unavailable")?;
+        let shadow_ctx = create_qwen_context(weights, self.compute_kind, &spec, false)?;
+        let mut candidate = Self {
+            state: self.state.clone(),
+            model_path: self.model_path.clone(),
+            config_toml: self.config_toml.clone(),
+            device: self.device,
+            compute_kind: self.compute_kind,
+            ctx: Some(shadow_ctx),
+            weights: None,
+            dataset: None,
+            lora_rank: self.lora_rank,
+            lora_alpha: self.lora_alpha,
+            lora_target_layers: self.lora_target_layers.clone(),
+            lora_target_modules: self.lora_target_modules.clone(),
+            dynamic_lora_configs: std::collections::BTreeMap::new(),
+            dynamic_lora_optimizer_configs: std::collections::BTreeMap::new(),
+            lr: self.lr,
+            beta1: self.beta1,
+            beta2: self.beta2,
+            eps: self.eps,
+            metrics: None,
+            last_loss: self.last_loss,
+            step: self.step,
+            context_spec: Some(spec),
+            pending_checkpoint_load: None,
+            _nccl_ep: self._nccl_ep,
+        };
+        let (step, loss) =
+            <Self as TrainingSession>::load_checkpoint_in_place(&mut candidate, path)?;
+        self.pending_checkpoint_load = Some(PendingCheckpointLoad {
+            transaction_id: transaction_id.to_string(),
+            source_path: path.to_string(),
+            ctx: candidate
+                .ctx
+                .take()
+                .context("checkpoint candidate lost its native context")?,
+            dynamic_lora_configs: std::mem::take(&mut candidate.dynamic_lora_configs),
+            dynamic_lora_optimizer_configs: std::mem::take(
+                &mut candidate.dynamic_lora_optimizer_configs,
+            ),
+            state: candidate.state.clone(),
+            step,
+            last_loss: loss,
+        });
+        Ok((step, loss))
+    }
+
+    pub fn commit_checkpoint_load(&mut self, transaction_id: &str) -> Result<(u64, f64)> {
+        let pending = self
+            .pending_checkpoint_load
+            .as_ref()
+            .with_context(|| format!("checkpoint transaction {transaction_id} is not prepared"))?;
+        if pending.transaction_id != transaction_id {
+            bail!(
+                "checkpoint transaction {} is prepared, not {transaction_id}",
+                pending.transaction_id
+            );
+        }
+        let pending = self
+            .pending_checkpoint_load
+            .take()
+            .context("validated checkpoint transaction disappeared")?;
+        self.ctx = Some(pending.ctx);
+        self.dynamic_lora_configs = pending.dynamic_lora_configs;
+        self.dynamic_lora_optimizer_configs = pending.dynamic_lora_optimizer_configs;
+        self.state = pending.state;
+        self.step = pending.step;
+        self.last_loss = pending.last_loss;
+        Ok((self.step, self.last_loss))
+    }
+
+    pub fn abort_checkpoint_load(&mut self, transaction_id: &str) -> Result<()> {
+        match self.pending_checkpoint_load.as_ref() {
+            Some(pending) if pending.transaction_id != transaction_id => bail!(
+                "checkpoint transaction {} is prepared, not {transaction_id}",
+                pending.transaction_id
+            ),
+            Some(_) => {
+                self.pending_checkpoint_load = None;
+                Ok(())
+            }
+            None => Ok(()),
+        }
+    }
+
+    fn load_checkpoint_transactional(&mut self, path: &str) -> Result<(u64, f64)> {
+        let transaction_id = format!(
+            "direct-{}-{}",
+            std::process::id(),
+            std::time::SystemTime::now()
+                .duration_since(std::time::UNIX_EPOCH)
+                .unwrap_or_default()
+                .as_nanos()
+        );
+        self.prepare_checkpoint_load(path, &transaction_id)?;
+        self.commit_checkpoint_load(&transaction_id)
+    }
+
+    pub fn export_distributed_adapter(
+        &self,
+        path: &str,
+        adapter_id: Option<i64>,
+        generation: &str,
+    ) -> Result<usize> {
+        let parallel = checkpoint::ParallelCheckpointManifest::from_env()?;
+        if parallel.world_size <= 1 {
+            return <Self as TrainingSession>::export_adapter(self, path, adapter_id);
+        }
+        if parallel.pipeline_model_parallel_size > 1 {
+            bail!(
+                "pipeline-parallel adapter export requires cross-stage tensor-name remapping; checkpoint save/restore remains supported"
+            );
+        }
+        let final_path = std::path::Path::new(path);
+        checkpoint::export_distributed_adapter_checkpoint(
+            final_path,
+            generation,
+            &parallel,
+            adapter_id,
+            |staging| {
+                <Self as TrainingSession>::save_checkpoint_with_generation(
+                    self,
+                    staging
+                        .to_str()
+                        .context("distributed export staging path is not UTF-8")?,
+                    Some(generation),
+                )
+                .map(|_| ())
+            },
+        )
+    }
 }
 
 impl TrainingSession for Qwen36Session {
@@ -200,57 +1246,35 @@ impl TrainingSession for Qwen36Session {
 
         // Load runtime config from model's config.json (no need to parse full TOML)
         let model_path_obj = std::path::Path::new(model_path);
-        let runtime_config =
-            rustrain_qwen3_6::config::read_qwen36_runtime_config(
-                model_path_obj,
-            )?;
+        let runtime_config = rustrain_qwen3_6::config::read_qwen36_runtime_config(model_path_obj)?;
 
-        // Build needed weight set — keys must match safetensors (with model prefix)
-        let n_layers = runtime_config.num_hidden_layers;
-        let wp = &runtime_config.weight_prefix;
-        let mut needed: std::collections::HashSet<String> = std::collections::HashSet::new();
-        needed.insert(format!("{wp}embed_tokens.weight"));
-        needed.insert(format!("{wp}norm.weight"));
-        if !runtime_config.tie_word_embeddings {
-            // lm_head.weight is always at top level (no model prefix), even for multimodal models
-            needed.insert("lm_head.weight".to_string());
-        }
-        for layer in 0..n_layers {
-            let p = format!("{wp}layers.{layer}");
-            needed.insert(format!("{p}.input_layernorm.weight"));
-            needed.insert(format!("{p}.post_attention_layernorm.weight"));
-            match runtime_config.layer_types[layer] {
-                rustrain_qwen3_6::config::LayerType::FullAttention => {
-                    for w in &["q_proj", "q_norm", "k_proj", "k_norm", "v_proj", "o_proj"] {
-                        needed.insert(format!("{p}.self_attn.{w}.weight"));
-                    }
-                }
-                rustrain_qwen3_6::config::LayerType::LinearAttention => {
-                    needed.insert(format!("{p}.linear_attn.in_proj_qkv.weight"));
-                    needed.insert(format!("{p}.linear_attn.in_proj_z.weight"));
-                    needed.insert(format!("{p}.linear_attn.in_proj_a.weight"));
-                    needed.insert(format!("{p}.linear_attn.in_proj_b.weight"));
-                    needed.insert(format!("{p}.linear_attn.A_log"));
-                    needed.insert(format!("{p}.linear_attn.dt_bias"));
-                    needed.insert(format!("{p}.linear_attn.conv1d.weight"));
-                    needed.insert(format!("{p}.linear_attn.norm.weight"));
-                    needed.insert(format!("{p}.linear_attn.out_proj.weight"));
-                }
-            }
-            if runtime_config.is_moe {
-                needed.insert(format!("{p}.mlp.gate.weight"));
-                needed.insert(format!("{p}.mlp.shared_expert_gate.weight"));
-                needed.insert(format!("{p}.mlp.shared_expert.gate_proj.weight"));
-                needed.insert(format!("{p}.mlp.shared_expert.up_proj.weight"));
-                needed.insert(format!("{p}.mlp.shared_expert.down_proj.weight"));
-                needed.insert(format!("{p}.mlp.experts.gate_up_proj"));
-                needed.insert(format!("{p}.mlp.experts.down_proj"));
-            } else {
-                needed.insert(format!("{p}.mlp.gate_proj.weight"));
-                needed.insert(format!("{p}.mlp.up_proj.weight"));
-                needed.insert(format!("{p}.mlp.down_proj.weight"));
-            }
+        let resolved = resolve_qwen_topology(&runtime_config)?;
+        let ResolvedQwenTopology {
+            ref topology,
+            global_rank,
+            world_size,
+            local_rank,
+            tp_rank,
+            cp_rank,
+            ep_rank: expert_rank,
+            dp_rank,
+            pp_rank,
+            tp_size,
+            cp_size,
+            ep_size,
+            dp_size,
+            pp_size,
+        } = resolved;
+        let stage = PipelineStageLayout::new(runtime_config.num_hidden_layers, pp_rank, pp_size)?;
+        if pp_size > 1 && runtime_config.mtp_num_hidden_layers > 0 {
+            bail!("pipeline-parallel Qwen server does not yet support MTP layers");
         }
+
+        // Resolve topology before disk IO so every worker reads only its
+        // frozen stage ownership set.
+        let wp = &runtime_config.weight_prefix;
+        let mut needed =
+            rustrain_qwen3_6::pipeline::stage_text_needed_weights(&runtime_config, &stage);
         // MTP weights
         if runtime_config.mtp_num_hidden_layers > 0 {
             for i in 0..runtime_config.mtp_num_hidden_layers {
@@ -286,47 +1310,237 @@ impl TrainingSession for Qwen36Session {
             model_path_obj,
             &needed,
         )?;
-
-        // ── Expert Parallel support ──
-        // Read EP params from env vars (set by launcher script)
-        let ep_rank = std::env::var("RANK").ok().and_then(|s| s.parse::<usize>().ok()).unwrap_or(0);
-        let ep_world_size = std::env::var("WORLD_SIZE").ok().and_then(|s| s.parse::<usize>().ok()).unwrap_or(1);
-        let local_rank = std::env::var("LOCAL_RANK").ok().and_then(|s| s.parse::<usize>().ok()).unwrap_or(0);
-        let is_ep = ep_world_size > 1 && runtime_config.is_moe;
+        let is_ep = runtime_config.is_moe && ep_size > 1;
+        let is_data_parallel = dp_size > 1;
+        let ep_a2a = std::env::var("QWEN36_EP_A2A")
+            .map(|value| !value.is_empty() && value != "0")
+            .unwrap_or(false);
+        let ep_a2a_sharded = std::env::var("QWEN36_EP_A2A_SHARDED")
+            .map(|value| !value.is_empty() && value != "0")
+            .unwrap_or(false);
+        validate_qwen_parallel_features(
+            runtime_config.is_moe,
+            tp_size,
+            ep_size,
+            ep_a2a,
+            ep_a2a_sharded,
+        )?;
+        let sequence_parallel = std::env::var("QWEN36_SEQUENCE_PARALLEL")
+            .map(|value| !value.is_empty() && value != "0")
+            .unwrap_or(false);
+        if sequence_parallel
+            && (runtime_config.is_moe
+                || runtime_config.has_vision
+                || tp_size != 2
+                || cp_size != 1
+                || ep_size != 1
+                || dp_size != 1
+                || pp_size != 1
+                || runtime_config.mtp_num_hidden_layers > 0
+                || runtime_config.router_aux_loss_coef != 0.0)
+        {
+            return Err(anyhow!(
+                "QWEN36_SEQUENCE_PARALLEL=1 currently requires dense text-only fixed-LoRA TP2 with CP=EP=DP=PP=1, MTP disabled, and router_aux_loss_coef=0"
+            ));
+        }
+        unsafe {
+            std::env::set_var("TP_SIZE", tp_size.to_string());
+            std::env::set_var("CP_SIZE", cp_size.to_string());
+            std::env::set_var("EP_SIZE", ep_size.to_string());
+            std::env::set_var("DP_SIZE", dp_size.to_string());
+            std::env::set_var("PP_SIZE", pp_size.to_string());
+            std::env::set_var("RUSTRAIN_TP_RANK", tp_rank.to_string());
+            std::env::set_var("RUSTRAIN_CP_RANK", cp_rank.to_string());
+            std::env::set_var("RUSTRAIN_EP_RANK", expert_rank.to_string());
+            std::env::set_var("RUSTRAIN_DP_RANK", dp_rank.to_string());
+            std::env::set_var("RUSTRAIN_PP_RANK", pp_rank.to_string());
+            std::env::set_var(
+                "RUSTRAIN_DATA_PARALLEL",
+                if is_data_parallel { "1" } else { "0" },
+            );
+        }
+        let base_tp_attention = tp_size > 1;
+        let base_tp_mlp = tp_size > 1;
+        let vocab_parallel = tp_size > 1;
+        if vocab_parallel
+            && (runtime_config.vocab_size <= 0 || runtime_config.vocab_size % tp_size as i64 != 0)
+        {
+            return Err(anyhow!(
+                "vocab_size={} must be divisible by TP_SIZE={tp_size} for vocabulary parallelism",
+                runtime_config.vocab_size
+            ));
+        }
+        if base_tp_attention {
+            if runtime_config.mtp_num_hidden_layers > 0 {
+                return Err(anyhow!(
+                    "frozen base TP currently requires MTP to be disabled"
+                ));
+            }
+            if runtime_config
+                .layer_types
+                .iter()
+                .any(|layer| *layer == rustrain_qwen3_6::config::LayerType::LinearAttention)
+                && (runtime_config.linear_num_key_heads <= 0
+                    || runtime_config.linear_num_value_heads <= 0
+                    || runtime_config.linear_num_value_heads % runtime_config.linear_num_key_heads
+                        != 0
+                    || runtime_config.linear_num_key_heads % tp_size as i64 != 0
+                    || runtime_config.linear_num_value_heads % tp_size as i64 != 0
+                    || runtime_config.linear_key_head_dim != 128
+                    || runtime_config.linear_value_head_dim != 128
+                    || runtime_config.linear_conv_kernel_dim <= 0)
+            {
+                return Err(anyhow!(
+                    "linear-attention TP requires k/v heads divisible by TP_SIZE with preserved value-head groups, 128-wide key/value heads, and a positive conv kernel: k_heads={}, v_heads={}, key_dim={}, value_dim={}, conv_kernel={}, tp={tp_size}",
+                    runtime_config.linear_num_key_heads,
+                    runtime_config.linear_num_value_heads,
+                    runtime_config.linear_key_head_dim,
+                    runtime_config.linear_value_head_dim,
+                    runtime_config.linear_conv_kernel_dim,
+                ));
+            }
+            if runtime_config.num_attention_heads <= 0
+                || runtime_config.num_attention_heads % tp_size as i64 != 0
+                || runtime_config.num_key_value_heads <= 0
+                || runtime_config.num_key_value_heads % tp_size as i64 != 0
+                || runtime_config.num_attention_heads % runtime_config.num_key_value_heads != 0
+            {
+                return Err(anyhow!(
+                    "full-attention heads (q={}, kv={}) must preserve GQA groups and be divisible by TP_SIZE={tp_size}",
+                    runtime_config.num_attention_heads,
+                    runtime_config.num_key_value_heads
+                ));
+            }
+            let rotary_dim =
+                (runtime_config.head_dim as f64 * runtime_config.partial_rotary_factor) as i64;
+            if runtime_config.head_dim <= 0
+                || rotary_dim < 0
+                || rotary_dim > runtime_config.head_dim
+                || rotary_dim % 2 != 0
+            {
+                return Err(anyhow!(
+                    "full-attention head_dim={} and partial_rotary_factor={} produce invalid rotary_dim={rotary_dim}",
+                    runtime_config.head_dim,
+                    runtime_config.partial_rotary_factor
+                ));
+            }
+        }
+        if base_tp_mlp {
+            let intermediates = if runtime_config.is_moe {
+                vec![
+                    ("routed expert", runtime_config.moe_intermediate_size),
+                    (
+                        "shared expert",
+                        runtime_config.shared_expert_intermediate_size,
+                    ),
+                ]
+            } else {
+                vec![("dense", runtime_config.intermediate_size)]
+            };
+            validate_tp_intermediate_sizes(tp_size, &intermediates)?;
+        }
 
         // Compute expert shard
         let (expert_start, expert_count) = if is_ep {
-            assert!(runtime_config.num_experts % ep_world_size == 0,
-                "num_experts {} not divisible by ep_world_size {}", runtime_config.num_experts, ep_world_size);
-            let epr = runtime_config.num_experts / ep_world_size;
-            (ep_rank * epr, epr)
+            assert!(
+                runtime_config.num_experts % ep_size == 0,
+                "num_experts {} not divisible by ep_size {}",
+                runtime_config.num_experts,
+                ep_size
+            );
+            let epr = runtime_config.num_experts / ep_size;
+            (expert_rank * epr, epr)
         } else {
             (0, runtime_config.num_experts)
         };
 
-        // Set CUDA device for EP
-        if is_ep {
+        // Set CUDA device for any torchrun worker. Dense Qwen workers use
+        // replicated weights and NCCL gradient all-reduce (LoRA-only DP).
+        if is_ep || is_data_parallel || tp_size > 1 {
             self.device = tch::Device::Cuda(local_rank);
         }
 
-        // Move to device — for EP, narrow expert tensors before GPU transfer
+        // Apply orthogonal EP and TP shards on CPU before moving weights to CUDA.
         let num_experts = runtime_config.num_experts as i64;
-        let mut weights: std::collections::BTreeMap<String, Tensor> = std::collections::BTreeMap::new();
+        let mut weights: std::collections::BTreeMap<String, Tensor> =
+            std::collections::BTreeMap::new();
         for (name, tensor) in raw_weights {
-            let needs_narrow = is_ep
-                && (name.contains(".mlp.experts.gate_up_proj") || name.contains(".mlp.experts.down_proj"));
-            if needs_narrow && tensor.size()[0] == num_experts {
-                let narrowed = tensor
-                    .narrow(0, expert_start as i64, expert_count as i64)
-                    .contiguous()
-                    .to_device(self.device)
-                    .to_kind(self.compute_kind);
-                weights.insert(name, narrowed);
+            let needs_expert_narrow = is_ep
+                && (name.contains(".mlp.experts.gate_up_proj")
+                    || name.contains(".mlp.experts.down_proj"));
+            let expert_shard = if needs_expert_narrow && tensor.size()[0] == num_experts {
+                Some(
+                    tensor
+                        .narrow(0, expert_start as i64, expert_count as i64)
+                        .contiguous(),
+                )
             } else {
-                let t = tensor.to_device(self.device);
-                let processed = t.to_kind(self.compute_kind);
-                weights.insert(name, processed);
-            }
+                None
+            };
+            let expert_or_full = expert_shard.as_ref().unwrap_or(&tensor);
+            let moe_tp_shard = if runtime_config.is_moe && base_tp_mlp {
+                rustrain_qwen3_6::kernel::shard_moe_mlp_weight_for_tp(
+                    &name,
+                    expert_or_full,
+                    tp_size,
+                    tp_rank,
+                )?
+            } else {
+                None
+            };
+            let vocab_shard = if expert_shard.is_none() && moe_tp_shard.is_none() && vocab_parallel
+            {
+                rustrain_qwen3_6::kernel::shard_vocab_weight_for_tp(
+                    &name,
+                    &tensor,
+                    runtime_config.vocab_size,
+                    tp_size,
+                    tp_rank,
+                )?
+            } else {
+                None
+            };
+            let local_shard = if moe_tp_shard.is_some() {
+                moe_tp_shard
+            } else if expert_shard.is_some() {
+                expert_shard
+            } else if vocab_shard.is_some() {
+                vocab_shard
+            } else if base_tp_attention {
+                let full_attention_shard =
+                    rustrain_qwen3_6::kernel::shard_full_attention_weight_for_tp(
+                        &name, &tensor, tp_size, tp_rank,
+                    )?;
+                let attention_shard = if full_attention_shard.is_some() {
+                    full_attention_shard
+                } else {
+                    rustrain_qwen3_6::kernel::shard_linear_attention_weight_for_tp(
+                        &name,
+                        &tensor,
+                        tp_size,
+                        tp_rank,
+                        runtime_config.linear_num_key_heads,
+                        runtime_config.linear_key_head_dim,
+                        runtime_config.linear_num_value_heads,
+                        runtime_config.linear_value_head_dim,
+                    )?
+                };
+                if attention_shard.is_some() || !base_tp_mlp {
+                    attention_shard
+                } else {
+                    rustrain_qwen3_6::kernel::shard_dense_mlp_weight_for_tp(
+                        &name, &tensor, tp_size, tp_rank,
+                    )?
+                }
+            } else {
+                None
+            };
+            let processed = local_shard
+                .as_ref()
+                .unwrap_or(&tensor)
+                .to_device(self.device)
+                .to_kind(self.compute_kind);
+            weights.insert(name, processed);
         }
 
         // Create C++ training context
@@ -336,41 +1550,108 @@ impl TrainingSession for Qwen36Session {
         } else {
             req.target_layers.clone()
         };
-        let lora_scaling = req.alpha as f64 / req.rank as f64;
-        let ctx = rustrain_qwen3_6::kernel::CppTrainingContext::new(
-            &weights,
+        let target_modules = req
+            .target_modules
+            .iter()
+            .map(|name| rustrain_qwen3_6::lora::Qwen36LoraTargetModule::parse(name))
+            .collect::<Result<Vec<_>>>()?;
+        rustrain_qwen3_6::lora::validate_lora_targets(
             &runtime_config,
-            self.compute_kind,
-            req.lr,
-            req.beta1,
-            req.beta2,
-            req.eps,
-            lora_scaling,
-            req.rank,
-            &all_layers,
-            expert_start,
-            expert_count,
+            &Qwen36LoraConfig {
+                rank: req.rank,
+                alpha: req.alpha,
+                target_layers: all_layers.clone(),
+                target_modules: target_modules.clone(),
+            },
         )?;
-
-        // Initialize NCCL communicator for EP — directly in C++
-        let nccl_ep = if is_ep {
-            let ret = ctx.init_nccl();
-            if ret != 0 {
-                return Err(anyhow!("C++ NCCL init failed (code {})", ret));
-            }
-            tracing::info!(ep_rank, ep_world_size, "NCCL communicator created in C++ for EP");
-            true
+        let (tp_color, cp_color, ep_color, dp_color, pp_color) = if world_size > 1 {
+            let tp_color = *topology
+                .tensor_group(global_rank)?
+                .iter()
+                .min()
+                .ok_or_else(|| anyhow!("empty TP process group"))?;
+            let cp_color = *topology
+                .context_group(global_rank)?
+                .iter()
+                .min()
+                .ok_or_else(|| anyhow!("empty CP process group"))?;
+            let ep_color = *topology
+                .expert_group(global_rank)?
+                .iter()
+                .min()
+                .ok_or_else(|| anyhow!("empty EP process group"))?;
+            let dp_color = *topology
+                .data_group(global_rank)?
+                .iter()
+                .min()
+                .ok_or_else(|| anyhow!("empty DP process group"))?;
+            let pp_color = *topology
+                .pipeline_group(global_rank)?
+                .iter()
+                .min()
+                .ok_or_else(|| anyhow!("empty PP process group"))?;
+            (tp_color, cp_color, ep_color, dp_color, pp_color)
         } else {
-            false
+            (0, 0, 0, 0, 0)
+        };
+        let context_spec = QwenContextSpec {
+            runtime_config: runtime_config.clone(),
+            stage,
+            init: req.clone(),
+            target_layers: all_layers.clone(),
+            target_modules: target_modules.clone(),
+            base_tp_attention,
+            base_tp_mlp,
+            vocab_parallel,
+            data_parallel: is_data_parallel,
+            expert_parallel: is_ep,
+            expert_start,
+            expert_count,
+            global_rank,
+            world_size,
+            tp_rank,
+            tp_size,
+            tp_color,
+            cp_rank,
+            cp_size,
+            cp_color,
+            ep_rank: expert_rank,
+            ep_size,
+            ep_color,
+            dp_rank,
+            dp_size,
+            dp_color,
+            pp_rank,
+            pp_size,
+            pp_color,
         };
+        let ctx = create_qwen_context(&weights, self.compute_kind, &context_spec, true)?;
+        let nccl_ep = world_size > 1;
+        if nccl_ep {
+            tracing::info!(
+                global_rank,
+                world_size,
+                data_parallel = is_data_parallel,
+                expert_parallel = is_ep,
+                tp_size,
+                "NCCL communicator created in C++ for Qwen parallel training"
+            );
+        }
 
         let count = ctx.lora_count() as usize;
         self.ctx = Some(ctx);
-        self.weights = Some(weights);  // Keep alive — C++ holds raw pointers
+        self.weights = Some(weights); // Keep alive — C++ holds raw pointers
+        self.context_spec = Some(context_spec);
+        self.pending_checkpoint_load = None;
         self._nccl_ep = nccl_ep;
         self.lora_rank = req.rank;
         self.lora_alpha = req.alpha;
+        self.lora_target_layers = all_layers;
+        self.lora_target_modules = target_modules;
         self.lr = req.lr;
+        self.beta1 = req.beta1;
+        self.beta2 = req.beta2;
+        self.eps = req.eps;
         self.state = SessionState::Ready {
             model_path: model_path.clone(),
         };
@@ -385,37 +1666,61 @@ impl TrainingSession for Qwen36Session {
             .ok_or_else(|| anyhow!("LoRA not initialized"))?;
 
         let loss = ctx.train_step(&input.input_ids, &input.target_mask, &input.attention_mask)?;
-        self.step += 1;
-        self.last_loss = loss;
-        self.state = SessionState::Training { step: self.step };
-
-        // Record metric
-        if let Some(ref metrics) = self.metrics {
-            let mem_gb = rustrain_train::metrics::gpu_memory_allocated_mb().map(|m| m / 1024.0).unwrap_or(0.0);
-            metrics.record_step(StepMetric {
-                step: self.step,
-                loss,
-                lr: self.lr,
-                mem_gb,
-                timestamp_unix: chrono::Utc::now().timestamp(),
-            });
-        }
-
-        Ok(TrainOutput { loss, step: self.step })
+        self.finish_train_step(loss, true)
     }
 
-    fn train_multi_lora(&mut self, input: TrainInput, n_total: i32, rank: i32) -> Result<TrainOutput> {
+    fn train_multi_lora(
+        &mut self,
+        input: TrainInput,
+        n_total: i32,
+        rank: i32,
+        adapter_ids: &[i64],
+    ) -> Result<TrainOutput> {
         let ctx = self
             .ctx
             .as_ref()
             .ok_or_else(|| anyhow!("LoRA not initialized"))?;
 
-        let loss = ctx.train_multi_lora(&input.input_ids, &input.target_mask, &input.attention_mask, n_total, rank)?;
-        self.step += 1;
-        self.last_loss = loss;
-        self.state = SessionState::Training { step: self.step };
-
-        Ok(TrainOutput { loss, step: self.step })
+        if n_total <= 0 {
+            return Err(anyhow!("n_total must be positive, got {n_total}"));
+        }
+        if !adapter_ids.is_empty() {
+            if adapter_ids.len() != n_total as usize {
+                return Err(anyhow!(
+                    "selected adapter count {} does not match n_total={n_total}",
+                    adapter_ids.len()
+                ));
+            }
+            if adapter_ids.iter().any(|id| *id <= 0) {
+                return Err(anyhow!("selected adapter IDs must be positive"));
+            }
+        }
+        if self
+            .context_spec
+            .as_ref()
+            .map(|spec| spec.pp_size > 1)
+            .unwrap_or(false)
+        {
+            return self.train_multi_lora_pipeline(&input, n_total, adapter_ids);
+        }
+        let loss = if adapter_ids.is_empty() {
+            ctx.train_multi_lora(
+                &input.input_ids,
+                &input.target_mask,
+                &input.attention_mask,
+                n_total,
+                rank,
+            )?
+        } else {
+            ctx.train_multi_lora_selected(
+                &input.input_ids,
+                &input.target_mask,
+                &input.attention_mask,
+                adapter_ids,
+                rank,
+            )?
+        };
+        self.finish_train_step(loss, false)
     }
 
     fn eval_step(&self, input: TrainInput) -> Result<EvalOutput> {
@@ -427,108 +1732,975 @@ impl TrainingSession for Qwen36Session {
         Ok(EvalOutput { loss })
     }
 
+    fn eval_multi_lora_host_i64(
+        &self,
+        input_ids: &[i64],
+        target_mask: &[i64],
+        attention_mask: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+        adapter_ids: &[i64],
+    ) -> Result<MultiLoraEvalOutput> {
+        Qwen36Session::eval_multi_lora_host_i64(
+            self,
+            input_ids,
+            target_mask,
+            attention_mask,
+            batch_size,
+            seq_len,
+            adapter_ids,
+        )
+    }
+
+    fn validate_dynamic_adapter_steps(
+        &self,
+        adapter_ids: &[i64],
+        expected_steps: &[u64],
+    ) -> Result<()> {
+        Qwen36Session::validate_dynamic_adapter_steps(self, adapter_ids, expected_steps)
+    }
+
     fn save_checkpoint(&self, path: &str) -> Result<(u64, f64)> {
+        self.save_checkpoint_with_generation(path, None)
+    }
+
+    fn save_checkpoint_with_generation(
+        &self,
+        path: &str,
+        checkpoint_generation: Option<&str>,
+    ) -> Result<(u64, f64)> {
         let ctx = self
             .ctx
             .as_ref()
             .ok_or_else(|| anyhow!("LoRA not initialized"))?;
 
-        let lora_count = ctx.lora_count();
-        let mut lora_a = Vec::new();
-        let mut lora_b = Vec::new();
-        for i in 0..lora_count {
-            if let (Some(a), Some(b)) = (ctx.get_lora_a(i as i64), ctx.get_lora_b(i as i64)) {
-                lora_a.push(a);
-                lora_b.push(b);
-            }
-        }
-
-        // Export Adam optimizer state
-        let (adam_m, adam_v) = ctx.export_optimizer_state()?;
-
-        checkpoint::save_checkpoint(
-            std::path::Path::new(path),
-            self.step,
-            self.last_loss,
-            self.model_path.as_deref().unwrap_or(""),
-            self.lora_rank,
-            self.lora_alpha,
-            &lora_a,
-            &lora_b,
-            &adam_m,
-            &adam_v,
-        )?;
+        let model_path = self
+            .model_path
+            .as_ref()
+            .ok_or_else(|| anyhow!("model path unavailable for checkpoint"))?;
+        let runtime_config =
+            rustrain_qwen3_6::config::read_qwen36_runtime_config(std::path::Path::new(model_path))?;
+        let parallel = checkpoint::ParallelCheckpointManifest::from_env()?;
+        let stage = self
+            .context_spec
+            .as_ref()
+            .context("LoRA context specification is unavailable")?
+            .stage
+            .clone();
+
+        let fixed_config = Qwen36LoraConfig {
+            rank: self.lora_rank,
+            alpha: self.lora_alpha,
+            target_layers: self.lora_target_layers.clone(),
+            target_modules: self.lora_target_modules.clone(),
+        };
+        let global_fixed_slots =
+            rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &fixed_config);
+        let fixed_slots = stage_lora_slots(&global_fixed_slots, &stage);
+        let lora_count = ctx.lora_count() as usize;
+        if fixed_slots.len() != lora_count {
+            bail!(
+                "fixed LoRA registry count {} does not match native slot count {lora_count}",
+                fixed_slots.len()
+            );
+        }
+        let (all_adam_m, all_adam_v) = ctx.export_optimizer_state()?;
+        let expected_optimizer_count = lora_count.saturating_mul(2);
+        if all_adam_m.len() != expected_optimizer_count
+            || all_adam_v.len() != expected_optimizer_count
+        {
+            bail!(
+                "fixed optimizer state count mismatch: m={}, v={}, expected={expected_optimizer_count}",
+                all_adam_m.len(),
+                all_adam_v.len()
+            );
+        }
+
+        // Distributed manifests compact inactive stage-local slots while their
+        // identities retain global layer and slot indices.
+        let saved_fixed_slots = fixed_slots
+            .iter()
+            .filter(|slot| parallel.world_size <= 1 || slot.active)
+            .collect::<Vec<_>>();
+        let saved_fixed_count = saved_fixed_slots.len();
+        let mut lora_a = Vec::with_capacity(saved_fixed_count);
+        let mut lora_b = Vec::with_capacity(saved_fixed_count);
+        let mut adam_m = Vec::with_capacity(saved_fixed_count.saturating_mul(2));
+        let mut adam_v = Vec::with_capacity(saved_fixed_count.saturating_mul(2));
+        let mut fixed_shard_layouts = Vec::with_capacity(saved_fixed_count);
+        let mut fixed_slot_identities = Vec::with_capacity(saved_fixed_count);
+        for slot in saved_fixed_slots {
+            lora_a.push(ctx.get_lora_a(slot.local_index as i64).with_context(|| {
+                format!(
+                    "fixed LoRA A is missing for native slot {}",
+                    slot.local_index
+                )
+            })?);
+            lora_b.push(ctx.get_lora_b(slot.local_index as i64).with_context(|| {
+                format!(
+                    "fixed LoRA B is missing for native slot {}",
+                    slot.local_index
+                )
+            })?);
+            let optimizer_index = slot.local_index.saturating_mul(2);
+            adam_m.push(all_adam_m[optimizer_index].shallow_clone());
+            adam_m.push(all_adam_m[optimizer_index + 1].shallow_clone());
+            adam_v.push(all_adam_v[optimizer_index].shallow_clone());
+            adam_v.push(all_adam_v[optimizer_index + 1].shallow_clone());
+            fixed_shard_layouts.push(checkpoint::lora_tp_shard_layout(
+                slot.module,
+                &runtime_config,
+            ));
+            fixed_slot_identities.push(checkpoint::LoraSlotIdentity {
+                index: slot.global_index,
+                layer: slot.layer,
+                module: slot.module.cpp_name().to_string(),
+            });
+        }
+
+        let supports_complete_dynamic_adam = ctx.supports_complete_dynamic_adam();
+        let mut dynamic_adapters = Vec::new();
+        if !self.dynamic_lora_configs.is_empty() {
+            for (&adapter_id, lora_config) in &self.dynamic_lora_configs {
+                let global_slots =
+                    rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, lora_config);
+                let slots = stage_lora_slots(&global_slots, &stage);
+                let shard_layouts = slots
+                    .iter()
+                    .filter(|slot| slot.active)
+                    .map(|slot| checkpoint::lora_tp_shard_layout(slot.module, &runtime_config))
+                    .collect::<Vec<_>>();
+                let optimizer_step = u64::try_from(ctx.get_adapter_step_count(adapter_id)?)
+                    .context("native dynamic adapter optimizer step is negative")?;
+                let mut dynamic_a = Vec::new();
+                let mut dynamic_b = Vec::new();
+                let mut dynamic_m = Vec::new();
+                let mut dynamic_v = Vec::new();
+                for slot in slots.iter().filter(|slot| slot.active) {
+                    let module = slot.module.cpp_name();
+                    dynamic_a.push(
+                        ctx.get_adapter_lora_tensor(
+                            adapter_id,
+                            slot.layer as i64,
+                            module,
+                            false,
+                        )
+                        .with_context(|| {
+                            format!(
+                                "dynamic LoRA A is missing: adapter={adapter_id} layer={} module={module}",
+                                slot.layer
+                            )
+                        })?,
+                    );
+                    dynamic_b.push(
+                        ctx.get_adapter_lora_tensor(
+                            adapter_id,
+                            slot.layer as i64,
+                            module,
+                            true,
+                        )
+                        .with_context(|| {
+                            format!(
+                                "dynamic LoRA B is missing: adapter={adapter_id} layer={} module={module}",
+                                slot.layer
+                            )
+                        })?,
+                    );
+                    if optimizer_step > 0 {
+                        // Keep one m/v entry for each A and B tensor, in slot order.
+                        dynamic_m.push(
+                            ctx.export_adapter_optimizer_tensor_cpu(
+                                adapter_id,
+                                slot.layer as i64,
+                                module,
+                                false,
+                                false,
+                            )
+                            .with_context(|| {
+                                format!(
+                                    "export dynamic LoRA m_a: adapter={adapter_id} layer={} module={module}",
+                                    slot.layer
+                                )
+                            })?
+                            .with_context(|| {
+                                format!(
+                                    "dynamic LoRA m_a is missing: adapter={adapter_id} layer={} module={module}",
+                                    slot.layer
+                                )
+                            })?,
+                        );
+                        dynamic_m.push(
+                            ctx.export_adapter_optimizer_tensor_cpu(
+                                adapter_id,
+                                slot.layer as i64,
+                                module,
+                                true,
+                                false,
+                            )
+                            .with_context(|| {
+                                format!(
+                                    "export dynamic LoRA m_b: adapter={adapter_id} layer={} module={module}",
+                                    slot.layer
+                                )
+                            })?
+                            .with_context(|| {
+                                format!(
+                                    "dynamic LoRA m_b is missing: adapter={adapter_id} layer={} module={module}",
+                                    slot.layer
+                                )
+                            })?,
+                        );
+                        dynamic_v.push(
+                            ctx.export_adapter_optimizer_tensor_cpu(
+                                adapter_id,
+                                slot.layer as i64,
+                                module,
+                                false,
+                                true,
+                            )
+                            .with_context(|| {
+                                format!(
+                                    "export dynamic LoRA v_a: adapter={adapter_id} layer={} module={module}",
+                                    slot.layer
+                                )
+                            })?
+                            .with_context(|| {
+                                format!(
+                                    "dynamic LoRA v_a is missing: adapter={adapter_id} layer={} module={module}",
+                                    slot.layer
+                                )
+                            })?,
+                        );
+                        dynamic_v.push(
+                            ctx.export_adapter_optimizer_tensor_cpu(
+                                adapter_id,
+                                slot.layer as i64,
+                                module,
+                                true,
+                                true,
+                            )
+                            .with_context(|| {
+                                format!(
+                                    "export dynamic LoRA v_b: adapter={adapter_id} layer={} module={module}",
+                                    slot.layer
+                                )
+                            })?
+                            .with_context(|| {
+                                format!(
+                                    "dynamic LoRA v_b is missing: adapter={adapter_id} layer={} module={module}",
+                                    slot.layer
+                                )
+                            })?,
+                        );
+                    }
+                }
+                let snapshot_step = u64::try_from(
+                    ctx.get_adapter_step_count(adapter_id)?,
+                )
+                .context("native dynamic adapter optimizer step is negative after export")?;
+                if snapshot_step != optimizer_step {
+                    bail!(
+                        "dynamic adapter {adapter_id} optimizer step changed during checkpoint export: {optimizer_step} -> {snapshot_step}"
+                    );
+                }
+                let optimizer = self
+                    .dynamic_lora_optimizer_configs
+                    .get(&adapter_id)
+                    .copied()
+                    .unwrap_or_else(|| self.default_dynamic_adam_config());
+                dynamic_adapters.push(checkpoint::DynamicAdapterCheckpoint {
+                    manifest: checkpoint::DynamicAdapterManifest {
+                        id: adapter_id,
+                        rank: lora_config.rank,
+                        alpha: lora_config.alpha,
+                        optimizer_step,
+                        optimizer_lr: Some(optimizer.lr),
+                        optimizer_beta1: supports_complete_dynamic_adam.then_some(optimizer.beta1),
+                        optimizer_beta2: supports_complete_dynamic_adam.then_some(optimizer.beta2),
+                        optimizer_eps: supports_complete_dynamic_adam.then_some(optimizer.eps),
+                        target_layers: lora_config.target_layers.clone(),
+                        target_modules: lora_config
+                            .target_modules
+                            .iter()
+                            .map(|module| module.cpp_name().to_string())
+                            .collect(),
+                        shard_layouts,
+                        slot_identities: slots
+                            .iter()
+                            .filter(|slot| slot.active)
+                            .map(|slot| checkpoint::LoraSlotIdentity {
+                                index: slot.global_index,
+                                layer: slot.layer,
+                                module: slot.module.cpp_name().to_string(),
+                            })
+                            .collect(),
+                        parameter_count: dynamic_a.len(),
+                        optimizer_count: dynamic_m.len(),
+                    },
+                    lora_a: dynamic_a,
+                    lora_b: dynamic_b,
+                    adam_m: dynamic_m,
+                    adam_v: dynamic_v,
+                });
+            }
+        }
+
+        let fixed_optimizer_step = u64::try_from(ctx.get_step_count())
+            .context("native fixed adapter optimizer step is negative")?;
+        if stage.pipeline_size > 1 {
+            let generation = checkpoint_generation
+                .filter(|generation| !generation.is_empty())
+                .context("pipeline-parallel checkpoint save requires a coordinated generation")?;
+            let stage_union = checkpoint::StageUnionCheckpointMetadata {
+                pipeline_stage: checkpoint::PipelineStageCheckpointManifest {
+                    pipeline_rank: stage.pipeline_rank,
+                    pipeline_size: stage.pipeline_size,
+                    global_num_layers: stage.global_num_layers,
+                    layer_start: stage.layer_range.start,
+                    layer_end: stage.layer_range.end,
+                },
+                fixed_target_layers: self.lora_target_layers.clone(),
+                fixed_target_modules: self
+                    .lora_target_modules
+                    .iter()
+                    .map(|module| module.cpp_name().to_string())
+                    .collect(),
+            };
+            checkpoint::save_stage_union_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+                std::path::Path::new(path),
+                self.step,
+                fixed_optimizer_step,
+                self.last_loss,
+                model_path,
+                self.lora_rank,
+                self.lora_alpha,
+                &lora_a,
+                &lora_b,
+                &adam_m,
+                &adam_v,
+                &dynamic_adapters,
+                &fixed_shard_layouts,
+                &fixed_slot_identities,
+                &parallel,
+                generation,
+                &stage_union,
+            )?;
+        } else {
+            match checkpoint_generation {
+                Some(generation) => {
+                    checkpoint::save_checkpoint_with_dynamic_and_fixed_step_for_topology_generation(
+                    std::path::Path::new(path),
+                    self.step,
+                    fixed_optimizer_step,
+                    self.last_loss,
+                    model_path,
+                    self.lora_rank,
+                    self.lora_alpha,
+                    &lora_a,
+                    &lora_b,
+                    &adam_m,
+                    &adam_v,
+                    &dynamic_adapters,
+                    &fixed_shard_layouts,
+                    &fixed_slot_identities,
+                    &parallel,
+                    Some(generation),
+                )?;
+                }
+                None => checkpoint::save_checkpoint_with_dynamic_and_fixed_step_for_topology(
+                    std::path::Path::new(path),
+                    self.step,
+                    fixed_optimizer_step,
+                    self.last_loss,
+                    model_path,
+                    self.lora_rank,
+                    self.lora_alpha,
+                    &lora_a,
+                    &lora_b,
+                    &adam_m,
+                    &adam_v,
+                    &dynamic_adapters,
+                    &fixed_shard_layouts,
+                    &fixed_slot_identities,
+                    &parallel,
+                )?,
+            }
+        }
 
         Ok((self.step, self.last_loss))
     }
 
     fn load_checkpoint(&mut self, path: &str) -> Result<(u64, f64)> {
-        let data = checkpoint::load_checkpoint(std::path::Path::new(path))?;
+        self.load_checkpoint_transactional(path)
+    }
+
+    fn load_checkpoint_in_place(&mut self, path: &str) -> Result<(u64, f64)> {
+        let parallel = checkpoint::ParallelCheckpointManifest::from_env()?;
+        let data = checkpoint::load_checkpoint_for_topology(std::path::Path::new(path), &parallel)?;
+        let stage = self
+            .context_spec
+            .as_ref()
+            .context("LoRA context specification is unavailable")?
+            .stage
+            .clone();
+        let model_path = self
+            .model_path
+            .as_ref()
+            .ok_or_else(|| anyhow!("model path unavailable for checkpoint restore"))?;
+        let current_model_path = std::fs::canonicalize(model_path)
+            .with_context(|| format!("canonicalize current base model path {model_path}"))?;
+        let checkpoint_model_path =
+            std::fs::canonicalize(&data.manifest.model_path).with_context(|| {
+                format!(
+                    "canonicalize checkpoint base model path {}",
+                    data.manifest.model_path
+                )
+            })?;
+        if checkpoint_model_path != current_model_path {
+            bail!(
+                "checkpoint base model {} does not match loaded model {}",
+                checkpoint_model_path.display(),
+                current_model_path.display()
+            );
+        }
+        if data.manifest.lora_rank != self.lora_rank
+            || data.manifest.lora_alpha.to_bits() != self.lora_alpha.to_bits()
+        {
+            bail!(
+                "checkpoint fixed LoRA rank/alpha {}/{} does not match session {}/{}",
+                data.manifest.lora_rank,
+                data.manifest.lora_alpha,
+                self.lora_rank,
+                self.lora_alpha
+            );
+        }
+        let runtime_config =
+            rustrain_qwen3_6::config::read_qwen36_runtime_config(std::path::Path::new(model_path))?;
+        let fixed_config = Qwen36LoraConfig {
+            rank: self.lora_rank,
+            alpha: self.lora_alpha,
+            target_layers: self.lora_target_layers.clone(),
+            target_modules: self.lora_target_modules.clone(),
+        };
+        if stage.pipeline_size > 1 {
+            let expected_stage = checkpoint::PipelineStageCheckpointManifest {
+                pipeline_rank: stage.pipeline_rank,
+                pipeline_size: stage.pipeline_size,
+                global_num_layers: stage.global_num_layers,
+                layer_start: stage.layer_range.start,
+                layer_end: stage.layer_range.end,
+            };
+            if data.manifest.pipeline_stage.as_ref() != Some(&expected_stage) {
+                bail!(
+                    "checkpoint pipeline stage metadata does not match the current runtime stage"
+                );
+            }
+            let expected_modules = self
+                .lora_target_modules
+                .iter()
+                .map(|module| module.cpp_name().to_string())
+                .collect::<Vec<_>>();
+            if data.manifest.fixed_target_layers != self.lora_target_layers
+                || data.manifest.fixed_target_modules != expected_modules
+            {
+                bail!("checkpoint global fixed LoRA target signature does not match the session");
+            }
+        }
+        let global_fixed_slots =
+            rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &fixed_config);
+        let fixed_slots = stage_lora_slots(&global_fixed_slots, &stage);
+        let expected_fixed_layouts = fixed_slots
+            .iter()
+            .filter(|slot| slot.active)
+            .map(|slot| checkpoint::lora_tp_shard_layout(slot.module, &runtime_config))
+            .collect::<Vec<_>>();
+        let expected_fixed_identities = fixed_slots
+            .iter()
+            .filter(|slot| slot.active)
+            .map(|slot| checkpoint::LoraSlotIdentity {
+                index: slot.global_index,
+                layer: slot.layer,
+                module: slot.module.cpp_name().to_string(),
+            })
+            .collect::<Vec<_>>();
+        if parallel.world_size > 1
+            || !data.manifest.fixed_shard_layouts.is_empty()
+            || !data.manifest.fixed_slot_identities.is_empty()
+        {
+            checkpoint::validate_fixed_tp_resume(
+                &data.manifest,
+                &expected_fixed_layouts,
+                &expected_fixed_identities,
+            )?;
+        }
+        if parallel.world_size > 1
+            || data.manifest.dynamic_adapters.iter().any(|adapter| {
+                !adapter.shard_layouts.is_empty() || !adapter.slot_identities.is_empty()
+            })
+        {
+            for dynamic in &data.dynamic_adapters {
+                let target_modules = dynamic
+                    .manifest
+                    .target_modules
+                    .iter()
+                    .map(|name| Qwen36LoraTargetModule::parse(name))
+                    .collect::<Result<Vec<_>>>()?;
+                let config = Qwen36LoraConfig {
+                    rank: dynamic.manifest.rank,
+                    alpha: dynamic.manifest.alpha,
+                    target_layers: dynamic.manifest.target_layers.clone(),
+                    target_modules,
+                };
+                let global_expected_slots =
+                    rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &config);
+                let expected_slots = stage_lora_slots(&global_expected_slots, &stage);
+                let expected_layouts = expected_slots
+                    .iter()
+                    .filter(|slot| slot.active)
+                    .map(|slot| checkpoint::lora_tp_shard_layout(slot.module, &runtime_config))
+                    .collect::<Vec<_>>();
+                let expected_identities = expected_slots
+                    .iter()
+                    .filter(|slot| slot.active)
+                    .map(|slot| checkpoint::LoraSlotIdentity {
+                        index: slot.global_index,
+                        layer: slot.layer,
+                        module: slot.module.cpp_name().to_string(),
+                    })
+                    .collect::<Vec<_>>();
+                checkpoint::validate_dynamic_tp_resume(
+                    &data.manifest,
+                    dynamic.manifest.id,
+                    &dynamic.manifest.shard_layouts,
+                    &expected_layouts,
+                    &expected_identities,
+                )?;
+            }
+        }
+        validate_dynamic_adapter_manifests(
+            data.dynamic_adapters
+                .iter()
+                .map(|dynamic| &dynamic.manifest),
+        )?;
+        for dynamic in &data.dynamic_adapters {
+            let manifest = &dynamic.manifest;
+            let target_modules = manifest
+                .target_modules
+                .iter()
+                .map(|name| Qwen36LoraTargetModule::parse(name))
+                .collect::<Result<Vec<_>>>()?;
+            let config = Qwen36LoraConfig {
+                rank: manifest.rank,
+                alpha: manifest.alpha,
+                target_layers: manifest.target_layers.clone(),
+                target_modules,
+            };
+            rustrain_qwen3_6::lora::validate_lora_targets(&runtime_config, &config)?;
+            let global_slots = rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &config);
+            let active_slots = stage_lora_slots(&global_slots, &stage)
+                .into_iter()
+                .filter(|slot| slot.active)
+                .count();
+            let full_optimizer_count = active_slots.saturating_mul(2);
+            let optimizer_count_valid = if dynamic.manifest.optimizer_step == 0 {
+                dynamic.manifest.optimizer_count == 0
+                    || dynamic.manifest.optimizer_count == full_optimizer_count
+            } else {
+                dynamic.manifest.optimizer_count == full_optimizer_count
+            };
+            if dynamic.lora_a.len() != active_slots
+                || dynamic.lora_b.len() != active_slots
+                || dynamic.adam_m.len() != dynamic.manifest.optimizer_count
+                || dynamic.adam_v.len() != dynamic.manifest.optimizer_count
+                || manifest.parameter_count != active_slots
+                || !optimizer_count_valid
+            {
+                bail!(
+                    "dynamic adapter {} tensor count does not match its runtime slot signature",
+                    manifest.id
+                );
+            }
+        }
+        if !data.dynamic_adapters.is_empty() {
+            if !self.dynamic_lora_configs.is_empty() {
+                bail!("cannot load dynamic LoRA checkpoint into a session with active adapters");
+            }
+            let ctx = self
+                .ctx
+                .as_ref()
+                .ok_or_else(|| anyhow!("LoRA not initialized"))?;
+            for dynamic in &data.dynamic_adapters {
+                let target_modules = dynamic
+                    .manifest
+                    .target_modules
+                    .iter()
+                    .map(|name| Qwen36LoraTargetModule::parse(name))
+                    .collect::<Result<Vec<_>>>()?;
+                let lora_config = Qwen36LoraConfig {
+                    rank: dynamic.manifest.rank,
+                    alpha: dynamic.manifest.alpha,
+                    target_layers: dynamic.manifest.target_layers.clone(),
+                    target_modules,
+                };
+                rustrain_qwen3_6::lora::validate_lora_targets(&runtime_config, &lora_config)?;
+                let layer_ids = lora_config
+                    .target_layers
+                    .iter()
+                    .map(|&layer| layer as i64)
+                    .collect::<Vec<_>>();
+                let module_csv = lora_config
+                    .target_modules
+                    .iter()
+                    .map(Qwen36LoraTargetModule::cpp_name)
+                    .collect::<Vec<_>>()
+                    .join(",");
+                let optimizer = DynamicAdamConfig {
+                    lr: dynamic.manifest.optimizer_lr.unwrap_or(self.lr),
+                    beta1: dynamic.manifest.optimizer_beta1.unwrap_or(self.beta1),
+                    beta2: dynamic.manifest.optimizer_beta2.unwrap_or(self.beta2),
+                    eps: dynamic.manifest.optimizer_eps.unwrap_or(self.eps),
+                };
+                let has_complete_override = dynamic.manifest.optimizer_beta1.is_some()
+                    || dynamic.manifest.optimizer_beta2.is_some()
+                    || dynamic.manifest.optimizer_eps.is_some();
+                let allocated_id = if has_complete_override {
+                    ctx.add_lora_for_restore_with_optimizer_config(
+                        lora_config.rank,
+                        lora_config.alpha,
+                        &layer_ids,
+                        &module_csv,
+                        optimizer,
+                    )?
+                } else if dynamic.manifest.optimizer_lr.is_some()
+                    && optimizer.lr.to_bits() != self.lr.to_bits()
+                {
+                    ctx.add_lora_for_restore_with_optimizer_lr(
+                        lora_config.rank,
+                        lora_config.alpha,
+                        &layer_ids,
+                        &module_csv,
+                        optimizer.lr,
+                    )?
+                } else {
+                    ctx.add_lora_for_restore(
+                        lora_config.rank,
+                        lora_config.alpha,
+                        &layer_ids,
+                        &module_csv,
+                    )?
+                };
+                if allocated_id != dynamic.manifest.id {
+                    if let Err(error) = ctx.set_adapter_id(allocated_id, dynamic.manifest.id) {
+                        let _ = ctx.remove_lora(allocated_id);
+                        return Err(error);
+                    }
+                }
+                let adapter_id = dynamic.manifest.id;
+                let optimizer_step = i64::try_from(dynamic.manifest.optimizer_step)
+                    .context("dynamic adapter optimizer step exceeds native range")?;
+                let load_result = (|| -> Result<()> {
+                    let global_slots =
+                        rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &lora_config);
+                    let slots = stage_lora_slots(&global_slots, &stage);
+                    let active_slots = slots
+                        .iter()
+                        .filter(|slot| slot.active)
+                        .collect::<Vec<_>>();
+                    let restore_optimizer = dynamic.manifest.optimizer_count != 0;
+                    for (slot_index, slot) in active_slots.iter().enumerate() {
+                        let module = slot.module.cpp_name();
+                        if slot_index >= dynamic.lora_a.len()
+                            || slot_index >= dynamic.lora_b.len()
+                        {
+                            bail!(
+                                "dynamic adapter {} tensor count mismatch",
+                                dynamic.manifest.id
+                            );
+                        }
+                        // A/B vectors are ordered exactly like active native slots.
+                        ctx.set_adapter_lora_tensor(
+                            adapter_id,
+                            slot.layer as i64,
+                            module,
+                            false,
+                            &dynamic.lora_a[slot_index],
+                        )?;
+                        ctx.set_adapter_lora_tensor(
+                            adapter_id,
+                            slot.layer as i64,
+                            module,
+                            true,
+                            &dynamic.lora_b[slot_index],
+                        )?;
+                    }
+                    if restore_optimizer || (active_slots.is_empty() && optimizer_step > 0) {
+                        let layers = active_slots
+                            .iter()
+                            .map(|slot| slot.layer as i64)
+                            .collect::<Vec<_>>();
+                        let modules = active_slots
+                            .iter()
+                            .map(|slot| slot.module.cpp_name())
+                            .collect::<Vec<_>>();
+                        ctx.import_adapter_optimizer_state_host(
+                            adapter_id,
+                            optimizer_step,
+                            &layers,
+                            &modules,
+                            &dynamic.adam_m,
+                            &dynamic.adam_v,
+                        )?;
+                    } else {
+                        ctx.set_adapter_step_count(adapter_id, optimizer_step)?;
+                    }
+                    Ok(())
+                })();
+                if let Err(error) = load_result {
+                    let _ = ctx.remove_lora(adapter_id);
+                    return Err(error);
+                }
+                self.dynamic_lora_configs.insert(adapter_id, lora_config);
+                self.dynamic_lora_optimizer_configs
+                    .insert(adapter_id, optimizer);
+            }
+        }
         // Import Adam optimizer state into C++ context
         if let Some(ctx) = &self.ctx {
-            if !data.adam_m.is_empty() && !data.adam_v.is_empty() {
-                ctx.import_optimizer_state(&data.adam_m, &data.adam_v)?;
-                tracing::info!(
-                    imported = data.adam_m.len(),
-                    "optimizer state imported"
+            let native_slot_count = ctx.lora_count() as usize;
+            if fixed_slots.len() != native_slot_count {
+                bail!(
+                    "fixed LoRA registry count {} does not match native slot count {native_slot_count}",
+                    fixed_slots.len()
+                );
+            }
+            let active_slot_indices = fixed_slots
+                .iter()
+                .filter(|slot| slot.active)
+                .map(|slot| slot.local_index)
+                .collect::<Vec<_>>();
+            let restore_slot_indices = checkpoint::fixed_restore_slot_indices(
+                data.lora_a.len(),
+                data.lora_b.len(),
+                &active_slot_indices,
+                native_slot_count,
+            )?;
+            for ((a, b), &slot_index) in data
+                .lora_a
+                .iter()
+                .zip(&data.lora_b)
+                .zip(&restore_slot_indices)
+            {
+                ctx.set_lora_tensor(slot_index as i64, false, a)?;
+                ctx.set_lora_tensor(slot_index as i64, true, b)?;
+            }
+            if data.adam_m.is_empty() != data.adam_v.is_empty() {
+                bail!(
+                    "checkpoint fixed optimizer m/v count mismatch: {}/{}",
+                    data.adam_m.len(),
+                    data.adam_v.len()
+                );
+            }
+            if data.adam_m.is_empty()
+                && !restore_slot_indices.is_empty()
+                && data.manifest.effective_fixed_optimizer_step() > 0
+            {
+                bail!(
+                    "checkpoint fixed optimizer step {} has no Adam state",
+                    data.manifest.effective_fixed_optimizer_step()
                 );
             }
+            if !data.adam_m.is_empty() {
+                let expected_saved_optimizer_count = restore_slot_indices.len().saturating_mul(2);
+                if data.adam_m.len() != expected_saved_optimizer_count
+                    || data.adam_v.len() != expected_saved_optimizer_count
+                {
+                    bail!(
+                        "checkpoint fixed optimizer count mismatch: m={}, v={}, expected={expected_saved_optimizer_count}",
+                        data.adam_m.len(),
+                        data.adam_v.len()
+                    );
+                }
+                let (mut all_adam_m, mut all_adam_v) = ctx.export_optimizer_state()?;
+                let expected_native_optimizer_count = native_slot_count.saturating_mul(2);
+                if all_adam_m.len() != expected_native_optimizer_count
+                    || all_adam_v.len() != expected_native_optimizer_count
+                {
+                    bail!(
+                        "native fixed optimizer count mismatch: m={}, v={}, expected={expected_native_optimizer_count}",
+                        all_adam_m.len(),
+                        all_adam_v.len()
+                    );
+                }
+                for (saved_slot, &native_slot) in restore_slot_indices.iter().enumerate() {
+                    let saved = saved_slot.saturating_mul(2);
+                    let native = native_slot.saturating_mul(2);
+                    all_adam_m[native] = data.adam_m[saved].shallow_clone();
+                    all_adam_m[native + 1] = data.adam_m[saved + 1].shallow_clone();
+                    all_adam_v[native] = data.adam_v[saved].shallow_clone();
+                    all_adam_v[native + 1] = data.adam_v[saved + 1].shallow_clone();
+                }
+                let imported = ctx.import_optimizer_state(&all_adam_m, &all_adam_v)?;
+                if imported != expected_native_optimizer_count as i64 {
+                    bail!(
+                        "native fixed optimizer import restored {imported} tensors, expected {expected_native_optimizer_count}"
+                    );
+                }
+                tracing::info!(imported, "optimizer state imported");
+            }
+            let native_step = i64::try_from(data.manifest.effective_fixed_optimizer_step())
+                .context("checkpoint step exceeds the native optimizer range")?;
+            ctx.set_step_count(native_step)?;
         }
         self.step = data.manifest.step;
         self.last_loss = data.manifest.loss;
+        self.state = SessionState::Paused { step: self.step };
         Ok((data.manifest.step, data.manifest.loss))
     }
 
-    fn export_adapter(&self, path: &str) -> Result<usize> {
+    fn export_adapter(&self, path: &str, adapter_id: Option<i64>) -> Result<usize> {
         let ctx = self
             .ctx
             .as_ref()
             .ok_or_else(|| anyhow!("LoRA not initialized"))?;
 
-        let count = ctx.lora_count() as usize;
-        let mut named_tensors: std::collections::BTreeMap<String, Tensor> =
-            std::collections::BTreeMap::new();
-        for i in 0..count {
-            if let (Some(a), Some(b)) = (ctx.get_lora_a(i as i64), ctx.get_lora_b(i as i64)) {
-                named_tensors.insert(
-                    format!("lora_a_{i}"),
-                    a.to_kind(Kind::Float).to_device(tch::Device::Cpu),
-                );
-                named_tensors.insert(
-                    format!("lora_b_{i}"),
-                    b.to_kind(Kind::Float).to_device(tch::Device::Cpu),
-                );
+        let model_path = self
+            .model_path
+            .as_ref()
+            .ok_or_else(|| anyhow!("model not loaded"))?;
+        let runtime_config =
+            rustrain_qwen3_6::config::read_qwen36_runtime_config(std::path::Path::new(model_path))?;
+        let adapter_id = adapter_id.unwrap_or(0);
+        let lora_config = if adapter_id == 0 {
+            Qwen36LoraConfig {
+                rank: self.lora_rank,
+                alpha: self.lora_alpha,
+                target_layers: self.lora_target_layers.clone(),
+                target_modules: self.lora_target_modules.clone(),
+            }
+        } else {
+            self.dynamic_lora_configs
+                .get(&adapter_id)
+                .cloned()
+                .ok_or_else(|| anyhow!("unknown dynamic LoRA adapter: {adapter_id}"))?
+        };
+        let slots = rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &lora_config);
+        let mut exported = Vec::with_capacity(slots.len());
+        for slot in slots {
+            if adapter_id == 0 {
+                let a = ctx
+                    .get_lora_a(slot.index as i64)
+                    .with_context(|| format!("native LoRA slot {} is missing A", slot.index))?;
+                let b = ctx
+                    .get_lora_b(slot.index as i64)
+                    .with_context(|| format!("native LoRA slot {} is missing B", slot.index))?;
+                exported.push((a, b));
+            } else if slot.active {
+                let module = slot.module.cpp_name();
+                let a = ctx
+                    .get_adapter_lora_tensor(adapter_id, slot.layer as i64, module, false)
+                    .with_context(|| {
+                        format!(
+                            "dynamic LoRA A is missing: adapter={adapter_id} layer={} module={module}",
+                            slot.layer
+                        )
+                    })?;
+                let b = ctx
+                    .get_adapter_lora_tensor(adapter_id, slot.layer as i64, module, true)
+                    .with_context(|| {
+                        format!(
+                            "dynamic LoRA B is missing: adapter={adapter_id} layer={} module={module}",
+                            slot.layer
+                        )
+                    })?;
+                exported.push((a, b));
+            } else {
+                let placeholder = Tensor::zeros([], (Kind::Float, Device::Cpu));
+                exported.push((placeholder.shallow_clone(), placeholder));
             }
         }
-
-        // Write safetensors
-        use std::io::Write;
-        let mut header = serde_json::Map::new();
-        let mut offset = 0u64;
-        let mut all_bytes: Vec<u8> = Vec::new();
-        for (name, t) in &named_tensors {
-            let t = t.contiguous().to_kind(Kind::Float);
-            let shape: Vec<i64> = t.size().iter().copied().collect();
-            let data: Vec<f32> = Vec::<f32>::try_from(&t.reshape([-1]))?;
-            let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
-            header.insert(
-                name.clone(),
-                serde_json::json!({"dtype":"F32","shape":shape,"data_offsets":[offset, offset + bytes.len() as u64]}),
-            );
-            offset += bytes.len() as u64;
-            all_bytes.extend_from_slice(&bytes);
-        }
-        let header_str = serde_json::to_string(&serde_json::Value::Object(header))?;
-        let file = std::fs::File::create(path)?;
-        let mut writer = std::io::BufWriter::new(file);
-        writer.write_all(&(header_str.len() as u64).to_le_bytes())?;
-        writer.write_all(header_str.as_bytes())?;
-        writer.write_all(&all_bytes)?;
-
+        let artifact = Qwen36AdapterArtifact::from_native_exports(
+            model_path,
+            "qwen3_hybrid_lora_sft",
+            Some(std::path::Path::new(model_path)),
+            &runtime_config,
+            &lora_config,
+            exported,
+        )?;
+        let count = artifact.tensors.len();
+        artifact.save(std::path::Path::new(path))?;
         tracing::info!(params = count, path, "adapter exported");
         Ok(count)
     }
 
+    fn import_adapter(&mut self, path: &str) -> Result<i64> {
+        let model_path = self
+            .model_path
+            .as_ref()
+            .ok_or_else(|| anyhow!("model not loaded"))?
+            .clone();
+        let runtime_config = rustrain_qwen3_6::config::read_qwen36_runtime_config(
+            std::path::Path::new(&model_path),
+        )?;
+        let artifact = Qwen36AdapterArtifact::load(std::path::Path::new(path))?;
+        let target_modules = artifact
+            .config
+            .target_modules
+            .iter()
+            .map(|name| Qwen36LoraTargetModule::parse(name))
+            .collect::<Result<Vec<_>>>()?;
+        let target_layers = artifact.config.target_layers.clone();
+        let lora_config = Qwen36LoraConfig {
+            rank: artifact.config.r,
+            alpha: artifact.config.lora_alpha,
+            target_layers: target_layers.clone(),
+            target_modules: target_modules.clone(),
+        };
+        rustrain_qwen3_6::lora::validate_lora_targets(&runtime_config, &lora_config)?;
+
+        let ctx = self
+            .ctx
+            .as_ref()
+            .ok_or_else(|| anyhow!("LoRA not initialized"))?;
+        let layer_ids = target_layers
+            .iter()
+            .map(|&layer| layer as i64)
+            .collect::<Vec<_>>();
+        let module_csv = target_modules
+            .iter()
+            .map(Qwen36LoraTargetModule::cpp_name)
+            .collect::<Vec<_>>()
+            .join(",");
+        let adapter_id =
+            ctx.add_lora(lora_config.rank, lora_config.alpha, &layer_ids, &module_csv)?;
+
+        let load_result = (|| -> Result<()> {
+            for (name, tensor) in &artifact.tensors {
+                let (layer, module, is_b) =
+                    rustrain_qwen3_6::lora::parse_adapter_tensor_name(&runtime_config, name)?;
+                if layer >= runtime_config.num_hidden_layers {
+                    bail!("adapter tensor layer {layer} is outside the model");
+                }
+                ctx.set_adapter_lora_tensor(
+                    adapter_id,
+                    layer as i64,
+                    module.cpp_name(),
+                    is_b,
+                    tensor,
+                )?;
+            }
+            Ok(())
+        })();
+        if let Err(error) = load_result {
+            let _ = ctx.remove_lora(adapter_id);
+            return Err(error);
+        }
+        self.dynamic_lora_configs.insert(adapter_id, lora_config);
+        self.dynamic_lora_optimizer_configs
+            .insert(adapter_id, self.default_dynamic_adam_config());
+        tracing::info!(adapter_id, path, "LoRA adapter imported");
+        Ok(adapter_id)
+    }
+
     fn status(&self) -> SessionStatus {
         let state = match &self.state {
             SessionState::Unloaded => "unloaded",
@@ -558,7 +2730,105 @@ impl TrainingSession for Qwen36Session {
             .ctx
             .as_ref()
             .ok_or_else(|| anyhow!("model not loaded — call load_model + init_lora first"))?;
-        let id = ctx.add_lora(req.rank, req.alpha, &req.target_layers, &req.target_modules)?;
+        let model_path = self
+            .model_path
+            .as_ref()
+            .ok_or_else(|| anyhow!("model not loaded"))?;
+        let runtime_config =
+            rustrain_qwen3_6::config::read_qwen36_runtime_config(std::path::Path::new(model_path))?;
+        let target_modules = if req.target_modules.trim().is_empty() {
+            let mut modules = std::collections::BTreeSet::new();
+            for layer_type in &runtime_config.layer_types {
+                match layer_type {
+                    rustrain_qwen3_6::config::LayerType::FullAttention => {
+                        modules.extend([
+                            Qwen36LoraTargetModule::QProj,
+                            Qwen36LoraTargetModule::KProj,
+                            Qwen36LoraTargetModule::VProj,
+                            Qwen36LoraTargetModule::OProj,
+                        ]);
+                    }
+                    rustrain_qwen3_6::config::LayerType::LinearAttention => {
+                        modules.extend([
+                            Qwen36LoraTargetModule::InProjQkv,
+                            Qwen36LoraTargetModule::InProjZ,
+                            Qwen36LoraTargetModule::InProjA,
+                            Qwen36LoraTargetModule::InProjB,
+                            Qwen36LoraTargetModule::OutProj,
+                        ]);
+                    }
+                }
+            }
+            if runtime_config.is_moe {
+                modules.extend([
+                    Qwen36LoraTargetModule::SharedGateProj,
+                    Qwen36LoraTargetModule::SharedUpProj,
+                    Qwen36LoraTargetModule::SharedDownProj,
+                    Qwen36LoraTargetModule::ExpertsGateUpProj,
+                    Qwen36LoraTargetModule::ExpertsDownProj,
+                ]);
+            } else {
+                modules.extend([
+                    Qwen36LoraTargetModule::GateProj,
+                    Qwen36LoraTargetModule::UpProj,
+                    Qwen36LoraTargetModule::DownProj,
+                ]);
+            }
+            modules.into_iter().collect()
+        } else {
+            req.target_modules
+                .split(',')
+                .filter(|name| !name.is_empty())
+                .map(Qwen36LoraTargetModule::parse)
+                .collect::<Result<Vec<_>>>()?
+        };
+        let config = Qwen36LoraConfig {
+            rank: req.rank,
+            alpha: req.alpha,
+            target_layers: req
+                .target_layers
+                .iter()
+                .map(|&layer| layer as usize)
+                .collect(),
+            target_modules: target_modules.clone(),
+        };
+        rustrain_qwen3_6::lora::validate_lora_targets(&runtime_config, &config)?;
+        let native_module_names = target_modules
+            .iter()
+            .map(Qwen36LoraTargetModule::cpp_name)
+            .collect::<Vec<_>>();
+        let module_csv = native_module_names.join(",");
+        let optimizer = DynamicAdamConfig {
+            lr: req.optimizer_lr.unwrap_or(self.lr),
+            beta1: req.optimizer_beta1.unwrap_or(self.beta1),
+            beta2: req.optimizer_beta2.unwrap_or(self.beta2),
+            eps: req.optimizer_eps.unwrap_or(self.eps),
+        };
+        let complete_override = req.optimizer_beta1.is_some()
+            || req.optimizer_beta2.is_some()
+            || req.optimizer_eps.is_some();
+        let id = if complete_override {
+            ctx.add_lora_with_optimizer_config(
+                req.rank,
+                req.alpha,
+                &req.target_layers,
+                &module_csv,
+                optimizer,
+            )?
+        } else {
+            match req.optimizer_lr {
+                Some(optimizer_lr) => ctx.add_lora_with_optimizer_lr(
+                    req.rank,
+                    req.alpha,
+                    &req.target_layers,
+                    &module_csv,
+                    optimizer_lr,
+                )?,
+                None => ctx.add_lora(req.rank, req.alpha, &req.target_layers, &module_csv)?,
+            }
+        };
+        self.dynamic_lora_configs.insert(id, config);
+        self.dynamic_lora_optimizer_configs.insert(id, optimizer);
         tracing::info!(adapter_id = id, rank = req.rank, "LoRA adapter added");
         Ok(id)
     }
@@ -570,6 +2840,8 @@ impl TrainingSession for Qwen36Session {
             .ok_or_else(|| anyhow!("model not loaded"))?;
         let removed = ctx.remove_lora(adapter_id)?;
         if removed {
+            self.dynamic_lora_configs.remove(&adapter_id);
+            self.dynamic_lora_optimizer_configs.remove(&adapter_id);
             tracing::info!(adapter_id, "LoRA adapter removed");
         }
         Ok(removed)
@@ -582,3 +2854,136 @@ impl TrainingSession for Qwen36Session {
             .unwrap_or_default()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{
+        dynamic_pipeline_schedule, validate_dynamic_adapter_manifests,
+        validate_qwen_parallel_features, validate_tp_intermediate_sizes,
+    };
+    use rustrain_qwen3_6::checkpoint::DynamicAdapterManifest;
+
+    fn dynamic_manifest(
+        id: i64,
+        rank: i64,
+        alpha: f64,
+        target_modules: &[&str],
+    ) -> DynamicAdapterManifest {
+        DynamicAdapterManifest {
+            id,
+            rank,
+            alpha,
+            optimizer_step: id as u64,
+            optimizer_lr: None,
+            optimizer_beta1: None,
+            optimizer_beta2: None,
+            optimizer_eps: None,
+            target_layers: vec![0],
+            target_modules: target_modules
+                .iter()
+                .map(|module| (*module).to_string())
+                .collect(),
+            shard_layouts: Vec::new(),
+            slot_identities: Vec::new(),
+            parameter_count: 1,
+            optimizer_count: 2,
+        }
+    }
+
+    #[test]
+    fn checkpoint_dynamic_manifests_allow_heterogeneous_signatures() {
+        let manifests = [
+            dynamic_manifest(1, 4, 8.0, &["q_proj", "v_proj"]),
+            dynamic_manifest(2, 16, 32.0, &["down_proj"]),
+        ];
+        validate_dynamic_adapter_manifests(manifests.iter()).unwrap();
+    }
+
+    #[test]
+    fn checkpoint_dynamic_manifests_reject_duplicate_ids() {
+        let manifests = [
+            dynamic_manifest(7, 4, 8.0, &["q_proj"]),
+            dynamic_manifest(7, 8, 16.0, &["down_proj"]),
+        ];
+        let error = validate_dynamic_adapter_manifests(manifests.iter()).unwrap_err();
+        assert!(error.to_string().contains("positive and unique"));
+    }
+
+    #[test]
+    fn checkpoint_dynamic_manifests_reject_invalid_optimizer_lr() {
+        let mut manifest = dynamic_manifest(1, 4, 8.0, &["q_proj"]);
+        manifest.optimizer_lr = Some(f64::NAN);
+        let error = validate_dynamic_adapter_manifests([&manifest]).unwrap_err();
+        assert!(error.to_string().contains("optimizer learning rate"));
+    }
+
+    #[test]
+    fn checkpoint_dynamic_manifests_reject_invalid_optimizer_scalars() {
+        let mut manifest = dynamic_manifest(1, 4, 8.0, &["q_proj"]);
+        manifest.optimizer_beta1 = Some(1.0);
+        let error = validate_dynamic_adapter_manifests([&manifest]).unwrap_err();
+        assert!(error.to_string().contains("optimizer beta1"));
+
+        manifest.optimizer_beta1 = Some(0.9);
+        manifest.optimizer_eps = Some(f64::NAN);
+        let error = validate_dynamic_adapter_manifests([&manifest]).unwrap_err();
+        assert!(error.to_string().contains("optimizer epsilon"));
+    }
+
+    #[test]
+    fn source_sharded_moe_tp_ep_is_supported() {
+        validate_qwen_parallel_features(true, 2, 2, true, true).unwrap();
+    }
+
+    #[test]
+    fn replicated_expert_tp_is_rejected() {
+        let error = validate_qwen_parallel_features(true, 2, 2, true, false).unwrap_err();
+        assert!(error.to_string().contains("replicated expert TP"));
+    }
+
+    #[test]
+    fn source_sharded_ep_requires_a2a() {
+        let error = validate_qwen_parallel_features(true, 2, 2, false, true).unwrap_err();
+        assert!(error.to_string().contains("requires QWEN36_EP_A2A=1"));
+    }
+
+    #[test]
+    fn source_sharding_requires_expert_parallelism() {
+        let error = validate_qwen_parallel_features(true, 2, 1, true, true).unwrap_err();
+        assert!(error.to_string().contains("expert-parallel training"));
+    }
+
+    #[test]
+    fn tp_mlp_accepts_divisible_routed_and_shared_intermediates() {
+        validate_tp_intermediate_sizes(2, &[("routed expert", 128), ("shared expert", 64)])
+            .unwrap();
+    }
+
+    #[test]
+    fn tp_mlp_rejects_non_divisible_shared_intermediate() {
+        let error =
+            validate_tp_intermediate_sizes(4, &[("routed expert", 128), ("shared expert", 66)])
+                .unwrap_err();
+        assert_eq!(
+            error.to_string(),
+            "shared expert intermediate_size=66 must be divisible by TP_SIZE=4"
+        );
+    }
+
+    #[test]
+    fn dynamic_pipeline_schedule_is_shape_safe_for_two_stages() {
+        assert_eq!(
+            dynamic_pipeline_schedule(0, 2, 3).unwrap(),
+            vec![
+                (Some(0), None),
+                (Some(1), Some(0)),
+                (Some(2), Some(1)),
+                (None, Some(2)),
+            ]
+        );
+        assert_eq!(
+            dynamic_pipeline_schedule(1, 2, 3).unwrap(),
+            vec![(Some(0), Some(0)), (Some(1), Some(1)), (Some(2), Some(2))]
+        );
+    }
+}
diff --git a/docs/agent/linear-attention.md b/docs/agent/linear-attention.md
index f5589545..eb940534 100644
--- a/docs/agent/linear-attention.md
+++ b/docs/agent/linear-attention.md
@@ -40,12 +40,83 @@ S = S + k_t ⊗ delta                 # state update
 out_t = S @ q_t                     # output
 ```
 
+## Tensor Parallel Layout
+
+GDN TP partitions K/V head groups while preserving `n_rep = V_heads / K_heads`:
+
+- `in_proj_qkv.weight` and `conv1d.weight`: slice Q, K, and V segments independently, then repack each rank as `[Q_local | K_local | V_local]`.
+- `in_proj_z.weight`: shard output rows by local V heads.
+- `in_proj_a.weight`, `in_proj_b.weight`, `A_log`, `dt_bias`: shard by local V-head rows.
+- `norm.weight`: replicate `[value_head_dim]`.
+- `out_proj.weight`: shard input columns matching local V heads, then all-reduce the local output.
+- Column-parallel LoRA replicates A and shards B; row-parallel LoRA shards A and replicates B. Replicated-factor gradients are summed once at the optimizer boundary.
+
+The v4 flat-QKV checkpoint manifest records three local-to-global row segments for Q, K, and V. Merge or reshard must apply those segments instead of treating the packed rank-local rows as one contiguous global slice.
+
+One `TpCopyToRegion` must wrap the shared input before the QKV/Z/A/B forks so backward sums their input-gradient contributions once. Do not all-reduce each fork separately.
+
+## Backward Stability
+
+The default CUDA backward never reconstructs earlier states by dividing by
+`g_exp`. Real Qwen weights can produce decay below `1e-8`, where subtracting the
+rank-one update and dividing by decay is numerically irreversible even though
+the forward and reference backward remain finite.
+
+Forward saves exact FP32 recurrent state at chunk boundaries. Backward reloads
+each chunk start, replays that chunk once, and retains exact `S_prev/S_t` pairs
+in a temporary per-batch-head workspace. The automatic checkpoint stride is
+`ceil(sqrt(sequence_length))`, which approximately minimizes the sum of saved
+forward boundaries and replay workspace. Set
+`QWEN36_GDN_STATE_CHECKPOINT_STRIDE=N` (`N >= 2`) to override it; values above
+the current sequence length are clamped to one full-sequence replay chunk. The
+two footprints are approximately:
+
+- saved states per live GDN layer: `BH * (ceil(S / N) + 1) * D_K * D_V * 4`
+- replay workspace: `BH * (N + 1) * D_K * D_V * 4`
+
+`QWEN36_GDN_INVERSE_BWD=1` explicitly selects the old low-memory inverse
+recurrence for diagnostics only. `QWEN36_GDN_CHUNKWISE_BWD=1` also retains its
+older inverse formulation and is not a correctness fallback for very small
+decay. Distributed registry consensus includes these settings so ranks cannot
+silently use different backward paths. The replay kernel remains one persistent
+block per `BH`; it is not FLA-style sequence-parallel chunking.
+
+## Right-Padding Fast Path
+
+Strict-right-padding masks are reduced once to device-resident `lengths[B]`.
+The persistent GDN forward and backward kernels stop at each sample's valid
+length and explicitly zero output and gradient tails. Chunked eval converts the
+global lengths to per-chunk offsets, and the batched LoRA path narrows lengths
+with the same adapter sub-batch as Q/K/V. Dense batches use an empty sentinel
+and retain the null-pointer fast path. Left padding and internal holes remain
+rejected until packed `cu_seqlens` boundaries are implemented.
+
+## Native GDN TP Verification
+
+Use a Python environment with ABI-compatible prebuilt PyTorch, CUDA, and NCCL, then run:
+
+```bash
+PYTHON=/path/to/python scripts/run_qwen36_native_gdn_tp.sh smoke
+PYTHON=/path/to/python scripts/run_qwen36_native_gdn_tp.sh bench-single
+PYTHON=/path/to/python scripts/run_qwen36_native_gdn_tp.sh bench-tp2
+```
+
+For checkpoint coverage, run the same smoke with
+`QWEN36_GDN_STATE_CHECKPOINT_STRIDE=2`. The synthetic training sequence has
+length 9, so this exercises five reverse chunks while retaining the distributed
+full-weight gradient and Adam-state oracle. The default smoke also injects decay
+below `1e-30` and directly compares every FP32 LoRA gradient accumulator from
+the stable replay backward with the ATen recurrence oracle.
+
+The script builds only the repository kernel and native harnesses. It discovers and links the prebuilt dependency files from the selected Python environment; missing headers or libraries are reported instead of building third-party dependencies.
+
 ## L2 Normalization
 
-- Transformers: `x * rsqrt(sum(x²) + eps)` (eps=1e-6, in denominator)
-- rustrain: `x / max(sqrt(sum(x²)), eps)` (eps=1e-6, as clamp)
+- Transformers/Megatron fused pre-GDN: `x * rsqrt(sum(x²) + eps)` (eps=1e-6)
+- rustrain C++ and Rust fallback: `x * rsqrt(sum(x²) + 1e-6)`
 
-Mathematically equivalent but numerically slightly different.
+The epsilon is inside the squared norm. Keeping this form is important for
+low-norm Q/K vectors; a post-sqrt clamp is not numerically equivalent.
 
 ## Diagnostic Dumps
 
diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md
new file mode 100644
index 00000000..c2b13ba3
--- /dev/null
+++ b/docs/plans/qwen-lora-megatron-progress.md
@@ -0,0 +1,176 @@
+---
+type: ProgressRecord
+title: Qwen LoRA Megatron runtime progress
+description: Verified milestones and open acceptance gaps for the Qwen native LoRA runtime.
+tags: [qwen3.5, qwen3.6, lora, progress]
+timestamp: 2026-07-17T00:00:00Z
+---
+
+# Source Spec
+
+[qwen-lora-megatron-spec.md](/docs/plans/qwen-lora-megatron-spec.md)
+
+# Current State
+
+The raw-binary HTTP multi-LoRA ingress is now implemented and verified by the local and H20 ABI1 server suites (55/55). The remaining ingress gap is asynchronous request/compute overlap.
+
+H20 ABI1 arbitrary-CP full-attention ring verification (2026-07-20): added opt-in `QWEN36_CP_FULL_ATTENTION_RING=1` with NCCL P2P KV rotation, FP32 online-softmax state, and a custom autograd backward that rotates FP32 dK/dV accumulators back to their KV owners. The existing CP2 KV all-gather remains the guarded oracle path; ring and gather are mutually exclusive and CP>2 requires ring. The CP smoke is parameterized for CP2/CP4: both CP2 and CP4 ring runs passed against a CP1 dense full-attention oracle with right padding, LoRA q/k/v/o gradients, Adam m/v and parameter updates, and rank-local flag divergence fail-closed. This is a memory-scalable correctness path, not a Megatron-throughput claim: communication is currently serialized with ATen matmul/softmax, with no zigzag partition, double-buffer overlap, fused tile kernel, or hybrid GDN CP>2 support.
+
+H20 ABI1 boundary verification (2026-07-20): `pp-train-smoke` exited successfully with fixed-LoRA 1F1B loss/gradient/parameter/Adam parity at both PP2 stages; the same run exercised fail-closed registry, selected-request, shape, and phase mismatch checks. `cp-gdn-smoke` exited successfully with CP2 GDN fixed and selected dynamic loss/tenant-loss differences of zero and bitwise-isolated unselected tenants. The same test intentionally rejects GDN `QWEN36_SEQ_CHUNK`, as sequence chunking is not yet a supported CP contract.
+
+H20 ABI1 full-attention CP2 bridge verification (2026-07-20): `cp-attention-smoke` passed a dense fixed-LoRA CP2 versus CP1 oracle with right padding crossing the CP boundary. Differentiable K/V all-gather, RoPE offsets, explicit global causal masking, train/eval, Adam state, and rank-local flag mismatch all passed. CP2-vs-CP1 eval/loss delta was `1.117e-3`; Adam `m/v` deltas were `5.595e-6/1.284e-9`. This remains the correctness oracle for the legacy gather path; the same smoke now also covers the ring path and CP4.
+
+The full-attention dynamic LoRA path now has an opt-in `QWEN36_FUSED_LORA_QKV_A=1` implementation. Compatible Q/K/V adapters share one concatenated A-side batched matmul and retain projection-specific B-side matmuls; heterogeneous layouts or ranks fall back to the established path, and the flag participates in the distributed runtime hash. H20 local and TP2 fixed/dynamic correctness smokes passed. Synthetic TP2 x EP2 A/B was effectively neutral: at `B=2,S=128,H=1024` p50 changed from `9.354` to `9.291 ms` (about `0.7%`), while repeated `B=8,S=512,H=2048` runs were noisy and showed no stable win. The flag therefore remains disabled by default.
+
+The MoE path now has an opt-in `QWEN36_MOE_SHARED_OVERLAP=1` stream/event implementation. It launches the shared expert on one context-owned nonblocking CUDA stream while routed sorting/A2A and local expert GEMMs remain on the compute stream, then waits only before the residual combine; the flag participates in the distributed runtime hash and no device-wide synchronize is used. H20 TP2 x EP2 fixed/dynamic smoke and failure-recovery cases passed. On the synthetic dynamic workload, small `B=2,S=128,H=1024` runs were noisy or slightly slower, and `B=8,S=512,H=2048` p50 changed from `53.195` to `53.576 ms` while allocator reserved memory increased by about `2.18 GiB/GPU`. It remains disabled by default until a workload with a stable overlap win and bounded workspace cost is demonstrated.
+
+MoE variable-split dispatch now packs token and expert IDs into one int64 metadata message, reducing each peer's forward dispatch from hidden/token/expert to hidden/metadata messages. `QWEN36_EP_A2A_PACKED_METADATA=0` retains the split fallback and the flag participates in the distributed runtime hash. H20 local and TP2 x EP2 smokes passed; a matched synthetic dynamic-MoE run measured split/packed p50 `14.155/14.051 ms` and p95 `16.849/16.016 ms`. This is a small, noise-scale result; host-visible variable-count planning and the lack of fused permutation/communication overlap remain the primary DeepEP gap.
+
+The EP A2A count path now materializes one CPU `int64` send/receive prefix plan immediately after the count all-gather and carries that plan through dispatch/combine autograd. Dispatch backward and combine forward/backward no longer repeat count `.to(CPU)`, `memcpy`, and prefix scans. H20 TP2 x EP2 dynamic matched A/B (`B=8,S=128,H=1024`, 10 warmup/50 iterations) measured baseline/prefix mean `13.974/13.977 ms`, p50 `13.793/14.008 ms`, p95 `15.243/14.736 ms`, with unchanged resident/reserved memory; this is a deterministic host-bookkeeping reduction, not a stable end-to-end speedup.
+
+Packed MoE now has an opt-in `QWEN36_MOE_FUSED_UNPERMUTE=1` weighted unpermute path. A CUDA inverse-order kernel directly reconstructs `[token, top_k]` output rows and its custom autograd backward emits both returned-activation and router-weight gradients, removing the packed path's `index_select`/temporary weighted tensor/`index_add` sequence. The flag is in the distributed runtime hash, supports BF16/FP16/FP32, and leaves legacy/non-packed dispatch unchanged. H20 TP2 x EP2 local and distributed smokes passed with fixed `adam_error=0`, resume parameter/m/v diff `0`, dynamic isolation diff `0`, heterogeneous-v2 `ok`, and recovery-poison `ok`; `cargo check -p rustrain-qwen3-6 --lib` also passed. Matched dynamic `B=8,S=128,H=1024` A/B (10 warmup/50 iterations) measured fused/fallback mean `14.063/14.154 ms`, p50 `13.963/14.132 ms`, p95 `15.248/15.223 ms`; the approximately `0.6%` mean difference is within run noise and resident memory was unchanged, so the flag remains disabled by default and does not close the fused permutation/DeepEP gap.
+
+Dynamic adapter optimizer residency now defaults to `QWEN36_LAZY_DYNAMIC_OPTIMIZER_STATE=1`: registration keeps live LoRA tensors and the FP32 gradient slab ready, while Adam `m/v` and transactional shadow tensors materialize only for a tenant with global tokens. Checkpoint step zero stores an implicit zero optimizer state (`optimizer_count=0`), while older full-state step-zero checkpoints remain loadable; state-only hydration does not allocate rollback shadows. H20 TP2 x EP2 dynamic-MoE A/B (`B=8,S=128,H=1024`, rank 16, 50 iterations) measured allocator reserved `1.035 GiB/GPU` for 8 registered/8 active and `2.918-3.148 GiB/GPU` for 1024 registered/8 active, versus the previous `9.46 GiB/GPU` full-residency result. Step p50 was `14.017 -> 14.058 ms`, within measurement noise. Native local/distributed smoke and both Rust crates' ABI1 `cargo check --lib` passed; lazy materialization failures join the existing collective fail-closed gate.
+
+Dynamic transactional Adam now defaults to `QWEN36_DYNAMIC_ADAM_SHADOW_POOL=1`. Tenant-owned `m/v` and gradient storage remain persistent optimizer state, but the six-tensor out-of-place undo buffer is leased from a context-level pool only for active tenants and is returned after the selected-v2 outer commit consensus or rollback. Exact layouts reuse an idle slot; a free slot is re-shaped before the pool grows, so mixed rank/target layouts remain bounded by the maximum simultaneous active pairs. H20 TP2 x EP2 with 1024 registered/8 active tenants and a full 128-step rotation measured old-path versus pooled `current/peak allocated` `8.439/8.813 -> 5.028/5.447 GiB`, `current/peak reserved` `9.029 -> 5.637 GiB`, and max resident `11.363 -> 7.971 GiB`; p50/p95 changed `21.179/32.224 -> 16.752/26.057 ms`. A second full 128-step rotation held allocated memory at `5.028 GiB` (reserved `5.879 GiB`), demonstrating no growth with historical tenant activation. H20 local and TP2 x EP2 distributed smokes, Adam/resume/heterogeneous/recovery checks, and the pool-disabled fallback all passed. Persistent optimizer state still grows with the number of tenants that have actually trained; bounded GPU residency for that state requires a later CPU/offload cache.
+
+The receive-side local-expert permutation now has an opt-in native CUDA counting-sort path, `QWEN36_MOE_FUSED_LOCAL_PERMUTE=1`. It replaces the local `sort/index_select/bincount/cumsum/index_add` sequence with a histogram, 1024-thread prefix scan, metadata scatter, and inverse row permutation; the permutation mapping is saved for autograd backward, while the existing path remains the default. The flag participates in the distributed runtime hash and requires no TE/DeepEP or JIT dependency. H20 TP2 x EP2 local/distributed smokes passed with fixed Adam error `0`, resume parameter/m/v diff `0`, heterogeneous-v2 `ok`, and recovery-poison `ok`; `cargo check -p rustrain-qwen3-6 --lib` passed. Matched dynamic B=8/S=128/H=1024/I=2048/rank16, packed A2A and fused weighted-unpermute A/B (10 warmup/50 iterations) measured fused/fallback mean `13.342/13.956 ms` and p50 `13.156/13.846 ms`; p95 was `15.808/14.792 ms` in this run, so the mean/p50 improvement is promising but tail latency is not yet stable. The flag remains disabled by default and does not provide Megatron/DeepEP communication overlap.
+
+The GDN persistent forward and checkpoint replay kernels now omit token-loop barriers for their thread-private recurrent columns; initialization and reverse-reduction barriers remain. H20 TP2 GDN smoke and chunkwise parity exited successfully with zero backward mismatches. Matched `B=2,S=512,H=2048,L=3,stride=32` replay p50 was `106.26 ms` baseline versus `106.73 ms` after the change (p90 `111.73` versus `107.26 ms`); default no-checkpoint p50 was `102.62` versus `102.34 ms`. This is a correctness-preserving micro-optimization with no stable multi-percent win, so it is kept as a low-risk codegen improvement rather than advertised as Megatron-level acceleration.
+
+The PP smoke now also asserts that `schedule != 0` and `num_chunks != 1` are rejected collectively. The native v1 window remains explicitly non-interleaved, single-chunk 1F1B; no chunk-local activation/gradient or PP-aware loader contract is implied by the ABI fields.
+
+The native host-i64 ingress also has an opt-in `QWEN36_HOST_PINNED_STAGING=1` path. Each context caches independent pinned input/target/attention buffers and enqueues non-blocking H2D after copying from the pageable IPC slab. On H20 TP2 x EP2 (`B=2,S=64,H=1024,E=8,I=2048,L=1`, 5 warmup/50 iterations), pageable staging measured `6.336 ms` p50 / `7.148 ms` p95 versus pinned `6.959 ms` / `7.108 ms`; losses matched (`5.411833`). The flag remains opt-in because it did not produce a stable end-to-end throughput win; true request/compute overlap remains open.
+
+Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP parity against a full-expert reference, variable-split EP A2A with fixed-LoRA data sharding, dense replicated-DP smoke with per-tenant token weighting, latent-rank and projection-aware LoRA TP, frozen full-attention/GDN/dense-MLP base-weight TP, routed/shared expert tensor parallelism, orthogonal dense TP x DP process groups, vocabulary-parallel embedding/LM-head/distributed cross entropy, composable LoRA-only TP x EP x expert-DP process groups with expert-local and dense-replicated gradient semantics, ABI27 five-axis TP/CP/EP/DP/PP communicator initialization and cached attach plus canonical contiguous PP stage ownership, topology-aware server source sharding, binary parent/worker tensor-slab transport, dynamic batch logical-step update, FP32 gradient storage/aggregation, standard Adam bias correction, independent fixed and per-tenant optimizer clocks, selected-tenant isolation, one padded heterogeneous selected-tenant forward/backward with one finalizer and a grouped rollback fallback, persistent out-of-place dynamic Adam shadow buffers, ABI31 host-authoritative Adam restore and opt-in cold-state paging, checkpoint v5 full-rank replica preflight and multi-axis layouts, shared CLI/server PEFT merge, same-topology CLI resume, two-phase distributed server save/load with shadow-context rollback, and 5D topology mapping.
+
+Not yet verified or implemented: interleaved/chunked PP scheduling and PP-aware loader/checkpoint placement, hybrid GDN CP>2 and multi-axis CP, DeepEP/TE/FLA prebuilt integration, asynchronous HTTP request/compute overlap, checkpoint immutable generation/`LATEST` replacement and power-loss durability, cross-topology checkpoint resharding, old-v3 attention checkpoint migration, content-fingerprint base-model identity, and matched Megatron throughput. Guarded dynamic MTP objective normalization is implemented for `TP=CP=EP=PP=1` (including DP count reduction), and the real Qwen3.6 MoE prediction-layer layout is now accepted on the guarded TP2 path; H20 TP2 MoE MTP matches a full-weight oracle. EP MTP loader ownership/communicator plumbing is present, but TPxEP and uneven source-token-count oracles remain pending; fixed-LoRA EP MTP stays fail-closed until its denominator is made global. The guarded TP2 dense sequence-parallel path now covers fixed and projection-aware dynamic LoRA, with native C++ pad-mask derivation, and is verified with H20 NCCL smoke; latent-rank dynamic targets remain rejected. Fixed-LoRA PP now supports fixed-shape, one-chunk non-interleaved 1F1B for any `PP_SIZE>=2`; H20 PP2 and PP3 smokes passed activation/gradient P2P, loss/gradient/Adam parity, late shape/phase failure consensus, and first-microbatch cross-stage contract rejection. Later local metadata failures complete fixed P2P with dummy contract-shaped activations, then finish performs a control all-reduce and clears gradients before rejecting the window. Dynamic adapter lifecycle operations now use a stage-invariant request hash across CP/PP, retain strict local-layout checks across TP/EP/DP, and reject registry mutation while a PP window is active. The opt-in dynamic-LoRA PP window reuses the same guarded 1F1B schedule, expands shared rows per tenant, restores per-tenant token numerators, and invokes the transactional dynamic finalizer; H20 PP2 smoke now verifies its two-tenant loss/Adam isolation. PP+CP MTP execution, auxiliary loss, and interleaving remain fail-closed. Sequence parallel remains restricted to `CP=EP=DP=PP=1`, MTP/aux-loss disabled, and projection-aware targets. The server now maps each global request to the same `(DP,EP)` source coordinate as the native CLI; TP peers share rows, sharded EP uses `DP*EP` sources, and replicated EP uses DP sources. Expert TP with expert sharding remains intentionally restricted to sharded A2A.
+
+The native Qwen3.6 path now implements the configured Switch-style `router_aux_loss_coef` for fixed-LoRA PP1/CP1 training. Router probabilities remain unchanged in the forward pass; the auxiliary gradient is attached through C++ autograd, padding tokens are excluded, and expert-token counts reduce over source-sharded EP and expert-DP while TP remains replicated. H20 ABI1 native smoke covered the checkpoint recompute path and observed a nonzero loss/Adam-state delta versus coefficient zero. Dynamic multi-LoRA and pipeline execution with a nonzero coefficient remain rejected explicitly. An opt-in dedicated NCCL stream was measured on the same H20 EP4 workload and removed after p50 regressed from `3.7746 ms` to `4.0496 ms`.
+
+EP count metadata now has an opt-in overlap path, `QWEN36_EP_A2A_COUNT_OVERLAP=1`. A dedicated non-blocking CUDA stream runs count all-gather and compact metadata D2H while the compute stream packs tokens; events preserve the allocation dependency and the flag participates in the runtime hash. H20 TP2xEP2 local/distributed correctness passed. Matched dynamic synthetic A/B (`B=8,S=128,H=1024,I=2048`) measured `13.625/13.807 ms` mean, while a larger `B=8,S=512,H=2048,I=4096` run was neutral (`52.758/52.772 ms`), so the flag remains disabled by default and is not reported as stable communication-compute overlap.
+
+Dynamic Adam checkpoint export/import now uses an ABI31 host snapshot. On the first save for a trained tenant, native C++ synchronizes the completed Adam step, copies the complete FP32 m/v state to a temporary CPU map, and atomically publishes it; four CPU exports then reuse that snapshot without re-materializing all device state. Restore imports the complete stage-local map in one transaction and keeps device m/v undefined; the first selected training call hydrates only that tenant's active pairs. Step-zero checkpoints retain the implicit `optimizer_count=0` representation. H20 local and TP2xEP2 distributed smokes covered step-zero rejection, CPU/contiguous/FP32 output, repeated-snapshot identity, partial-import rollback, zero resident state after restore, first-use hydration, and resume parity. This is bounded host-resident restore/save behavior, not full optimizer offload, eviction, or paging.
+
+The H20 prebuilt dependency audit found no compatible Transformer Engine, DeepEP, FLA, or standalone flash-attn package for the project ABI1 PyTorch 2.12.1/cu130 environment. PyTorch bundled Flash-SDPA is usable; no dependency was installed and no runtime JIT path was introduced.
+
+Dynamic tenants now also have a native selected-eval ABI, Rust host-i64 binding, and `POST /v1/sessions/{id}/eval_multi_lora` in both the regular and EP HTTP routers. Evaluation scopes one adapter registry entry at a time and returns per-tenant losses; H20 single-rank smoke confirmed two registered tenants evaluate successfully without advancing optimizer clocks. The route is intentionally read-only, rejects empty/non-positive selections at the API boundary, and preserves native collective validation for unknown IDs.
+
+The guarded sequence-parallel path uses Megatron-style sequence ownership: embedding slices the full sequence, column-parallel projections all-gather before GEMM, row-parallel projections reduce-scatter their outputs, and vocabulary loss gathers hidden states before cross-entropy. The native topology validator rejects MoE, latent-rank dynamic LoRA, MTP, router auxiliary loss, and any CP/EP/DP/PP combination while this opt-in is enabled. H20 TP2 smoke passed fixed and three-tenant selected dynamic LoRA on both ranks with unequal tenant token counts; packed and per-tensor LoRA synchronization both passed, and the native C++ pad-token mask path was exercised with null Rust attention masks. The distributed-vs-single-rank dynamic oracle reported loss differences below `1e-2`, parameter differences below `2e-3`, and preserved the unselected tenant bitwise.
+
+Dynamic selected training now quarantines a context when registry restoration, gradient cleanup, transactional rollback, or registry restore fails during secondary recovery. The next dynamic request performs a TP/EP/DP health consensus before touching the tenant registry, propagates quarantine to the full worker group, and fails closed until the workers are recreated. The optional ABI28 health symbol is wired through Rust; EP workers return a terminal result and the parent coordinator immediately becomes unavailable and reaps the group. H20 world4 `TP2 x EP2` passed a rank0-only poison injection on all ranks without a collective hang, while the normal injected Adam failure still recovers with a healthy context.
+
+The follow-up H20 ABI29 smoke verifies TP2 dense sequence parallel with selected dynamic LoRA: three registered tenants, two selected rows with unequal token counts, full-attention Q/K/V/O plus gate/up/down projection-aware shards, and a single-rank full-weight oracle. Both packed and per-tensor LoRA gradient synchronization pass; selected parameter and Adam-state differences stay within BF16 reduction-order bounds, while the unselected tenant remains bitwise unchanged. The native pad-token mask is now derived in C++ when Rust passes a null attention mask, and the same path is covered by fixed micro-step, dynamic selected training, and PP tick tests. H20 local, TP-attention, CP/GDN, and PP2 regression smokes remain green; the qwen crate also passes `cargo check --lib` in the ABI1 H20 venv.
+
+Dynamic PP now keeps heterogeneous padding enabled for the complete selected 1F1B window and restores the previous context mode on reset. The PP2 H20 smoke uses selected rank-2/rank-3 tenants plus an unselected tenant, verifies loss/Adam isolation, and injects a second-microbatch metadata failure; the fallback keeps the selected-row batch size and both stages fail uniformly at finish without a P2P count mismatch.
+
+ABI28 selected-v2 preflight now packs health, request validity, report capability/buffer, and accumulation-clear flags into one `int32[6]` MIN collective per EP/DP/TP axis. H20 TP2 x EP2 heterogeneous B=8/S=128 A/B runs against the legacy six-call implementation measured p50 old/packed pairs of `13.394/13.256`, `13.219/13.153`, and `13.638/13.199 ms`; the median improved about 1.5%, and a matched Nsight Systems trace reduced u32 all-reduce instances from `1032` to `752`. World4 heterogeneous, rank-local failure, report-mode mismatch, rollback, and poison smokes remained green.
+
+ABI29 GDN CP2 now has an opt-in fused QKV+Z/A/B sequence-to-head exchange. The peer-major payload is `[Q|K|V|Z|A|B]`; forward and autograd reverse exchange each use one NCCL SendRecv group, and the flag participates in the cross-rank runtime hash. Fixed and dynamic CP2 full-reference smokes passed with the flag disabled and enabled, including unequal-token selected tenants, unselected bitwise isolation, and rank-local flag mismatch rejection. On H20 `B=2,S=512,H=2048,L=3,rank=8,warmup=5,iters=30`, rank0 p50/p90 was legacy/fused `103.811/105.456` vs `104.011/107.565 ms`; at B=8 it was `227.745/231.473` vs `227.952/236.727 ms`. A minimal Nsight trace reduced SendRecv kernels `36 -> 24`; end-to-end latency remained neutral to slightly worse, so the optimization remains opt-in and the default legacy path preserves projection ordering and activation lifetime.
+
+Dynamic MTP normalization is enabled for the dense single-axis topology `TP=CP=EP=PP=1` (with `DP>=1`) and for a guarded TP2 full-attention prediction-layer path; router auxiliary loss remains disabled. Main and MTP token counts are reduced together over DP, checked for replica equality over TP, and sharded-EP counts now reduce before DP while replicated EP counts are equality-checked. MTP hidden gradients/reports are converted to the common main-loss denominator. The original H20 ABI1 `mtp-dynamic-smoke` and `mtp-dp-smoke` remain exact against their single-tenant oracles. The H20 `mtp-tp-smoke` shards prediction-layer Q/K/V/O and dense gate/up/down weights, keeps embedding/LM-head replicated, and matches a full-weight single-rank oracle on both ranks (`loss_diff=3.53e-3`, parameter diff `1.87e-9`, Adam m/v diffs `2.14e-5/3.23e-9`, unselected tenant bitwise unchanged); the same oracle passes with fused QKV and FC1 caches enabled. The same TP2 oracle now also uses a real 4-expert routed/shared prediction layer (`loss_diff=1.50e-3`, parameter diff `3.05e-5`, unselected tenant bitwise unchanged), including fused QKV/FC1 caches. Vocabulary-parallel MTP also passes on TP2 (`loss_diff=3.62e-3`, parameter diff `2.01e-3`), and the TP2xEP2 uneven-source oracle passes with main/MTP token counts `7/5`, loss diff `4.38e-4`, and parameter diff `2.01e-3`. Dynamic multi-LoRA chunk validity/loss checks now remain on device until one final logical-step consensus, reducing host fences from per chunk to one per step while preserving fail-closed Adam/clock commit. TP2 MTP requires `TP_SIZE=2`, both base TP flags, no sequence parallelism, and `CP=PP=1`; PP/CP MTP and full-model long-run verification remain pending. Fixed-LoRA EP MTP stays fail-closed until its denominator is made global.
+
+# Durable Milestones
+
+- `4e242ee`: added Megatron-style 5D topology contract and launcher normalization; `cargo test -p rustrain-parallel --lib` passed 14/14.
+- ABI26 working tree: C++ and Rust now initialize TP/CP/EP/DP/PP communicators in a fixed collective order, require cross-rank axis-size consensus, atomically publish generation-scoped file rendezvous state, cache/attach all five axes, and reject mismatched context coordinates. Launcher jobs use the shared `RUSTRAIN_LAUNCH_OUTPUT_DIR`; direct jobs require a unique `RUSTRAIN_NCCL_RUN_ID` and may set `RUSTRAIN_NCCL_SYNC_DIR` for shared multi-node storage. H20 CP2 x PP2 world4 propagated rank max `3.0` through groups `[0,1]/[2,3]` and `[0,2]/[1,3]`; cached attach passed, invalid coordinates failed, and train/eval remained fail-closed. TP2 x EP2 and TP2 GDN regressions retained identical heterogeneous loss and `adam_error=0`.
+- ABI27 working tree: Rust and C++ share a canonical contiguous PP layout with global layer IDs, stage-local pointer/config arrays, first-stage embedding ownership, last-stage norm/LM-head ownership, tied-vocabulary duplication only at the two boundaries, and strict range/weight-count/null-pointer validation. Fixed LoRA target IDs are mapped global-to-local while exported native names retain their global layer prefix. H20 CP2 x PP2 world4 created `[0,1)/2` and `[1,2)/2` contexts with seven local slots each; global layer 0 was active only on its owner stage, cached communicator attach and max propagation passed on all ranks, and unsupported PP/CP execution remained fail-closed. The full single-rank native smoke retained fixed/dynamic/heterogeneous multi-LoRA, GDN, dense, and standard-Adam oracles.
+- PP generalization working tree: fixed-LoRA model execution now runs fixed-shape, one-chunk non-interleaved 1F1B for any `PP_SIZE>=2`. PP control collectives use a duplicate communicator; only the first microbatch establishes the cross-stage shape/dtype contract, because per-tick collectives would lock stages together and deadlock against asynchronous activation/gradient P2P. Later metadata failures are converted to contract-shaped dummy work so all ranks consume the same P2P schedule; finish performs a control consensus and clears gradients before rejecting any window with a local error. H20 PP2 and PP3 `pp-train-smoke` passed normal loss/gradient/parameter/Adam parity, late shape+phase failure, and the divergent-target negative case.
+- Dynamic PP lifecycle hardening: canonical adapter registration no longer hashes stage-local `num_layers`, `global_layer_start`, or tensor layouts on the CP/PP request axes; those layouts remain mandatory on TP/EP/DP peers that share ownership. Stage-local failures still enter the CP/PP request consensus, health and mutation-clear consensus include PP, and add/rename/remove are rejected for an active pipeline window. H20 `CP2 x PP2` passed a one-replica target mismatch without hanging, then registered/renamed/removed the valid global-layer request on all four ranks. H20 PP2 rejected add-adapter on both stages during an active window and then completed the original 1F1B window with unchanged loss/gradient/Adam parity.
+- `afcf091`: restored native C++ Adam step on checkpoint load and added ABI8 smoke assertions; server checkpoint tests passed 2/2.
+- ABI9 working tree: selected adapter IDs flow through HTTP, IPC, Rust, and C++; each dynamic tenant owns its Adam clock, checkpoint metadata preserves it, and failed selection restores the complete registry.
+- ABI10 working tree: TP-only mode shards each fixed and dynamic adapter's latent rank, all-reduces only activation-level LoRA deltas on a split TP communicator, and keeps replicated MoE layers off the EP communicator.
+- `5aa4ea4`: ABI11 stores fixed and dynamic LoRA leaf gradients in FP32 accumulators between micro-batches; fused Adam consumes those accumulators directly, fixed windows use token-weighted numerator/denominator reduction, and explicit abort/failure cleanup clears the pending window. This is FP32 storage/aggregation; the autograd leaf backward remains BF16-typed.
+- `ed25daa`: TP rank-aware checkpoint format v3 writes `rank-xxxxx/` manifests and adapter/optimizer shards with topology, global/local shape, latent partition axis, offset, and replica identity metadata. v1/v2 single-rank loading remains compatible; v3 requires the same topology and rank.
+- `ff382a0`: native ABI rejects unsupported TP+DP/EP mixtures before training. Dynamic per-tenant DP batches now use an all-reduced per-adapter token-count vector and weighted FP32 LoRA gradient reduction, including grouped expert LoRA.
+- `76eace6`: direct native context creation also rejects PP/CP sizes greater than one, so unsupported pipeline/context parallelism cannot silently fall back to replicated single-stage execution.
+- `b7050ea`: fixed standard Adam bias correction, separated pure-DP gradient NCCL from MoE activation collectives, made dynamic-MoE DP token weighting numerical, scoped n_max rendezvous by invocation with atomic publication, and added full-expert EP parity plus Adam oracles.
+- `bfa3d8e`: added a gated variable-split NCCL dispatch/inverse-combine prototype with differentiable forward/backward collectives; default legacy EP remains unchanged.
+- Working tree after `bfa3d8e`: `QWEN36_EP_A2A_SHARDED=1` disjoins EP input rows, all-reduces global fixed-LoRA token numerators, and keeps grouped expert LoRA local after A2A. Native dynamic multi-LoRA now reuses the transported source row for tenant mapping, all-reduces per-tenant counts, skips zero-global-token tenants, and rejects ordinary single-adapter entry points while a dynamic registry is active.
+- H20 `123.57.26.97:28004`: ABI8 native smoke passed grouped/fallback parity, GDN, dense/MoE LoRA, dynamic adapters, and step setter validation.
+- H20 `123.57.26.97:28004`: ABI9 native smoke passed selected-tenant training with a positive selected update, exactly zero unselected update, independent clocks (`2` vs `1`), and registry preservation after an unknown ID.
+- H20 `123.57.26.97:28004`: ABI10 two-rank TP native smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, and selected-tenant isolation. Rank-local LoRA tensors used distinct rank `4` slices for global rank `8`; losses matched and both shards had positive updates.
+- H20 `123.57.26.97:28004`: ABI11 single-rank smoke passed FP32 accumulator dtype, two-micro accumulation, NaN/explicit abort cleanup, successful step commit, standard Adam parameter oracle, TP/DP/EP and PP/CP topology rejection, and dynamic multi-LoRA. ABI11 two-rank TP smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, selected-tenant isolation, invocation-scoped n_max rendezvous, and the new guards (`rank_statuses=0,0`).
+- H20 `123.57.26.97:28004`: ABI11 DP2 smoke passed with per-tenant masks `[1,3]` and `[3,1]`: weighted m relative error `2.43e-8`, grouped-expert error `2.27e-8`, v error `7.33e-8`, BF16 Adam delta error `0`, and nonzero gap `7.96e-3` versus the old equal-count formula (`rank_statuses=0,0`).
+- H20 `123.57.26.97:28004`: ABI11 EP2 full-expert parity smoke passed with distinct rank-local base and LoRA expert slices; loss, A/B updates, m/v, and standard Adam first-step oracle all matched the rank-local full-expert reference (`rank_statuses=0,0`).
+- H20 `123.57.26.97:28004`: gated A2A EP2 smoke passed in legacy, replicated-source, and sharded-source modes (`rank_statuses=0,0`). Sharded mode used token counts `[1,3]`; weighted global loss differed from the full-batch reference by `9.6e-7`, expert m/v maxima were `1.22e-5` / `3.92e-9`, and every standard Adam oracle was zero. One distributed BF16 parameter landed in the adjacent rounding bin (`9.61e-4`), bounded separately by the smoke.
+- Target runtime probe: PyTorch 2.5.1+cu121, ABI0; Transformer Engine, flash-attn, DeepEP, Triton, and DeepSpeed are not importable.
+- H20 `123.57.26.97:28004`: `native_ep_bench.cpp` fresh ABI0 benchmark (`seq=128, hidden=256, experts=8, intermediate=256, warmup=2, iters=10`) passed legacy and sharded A2A with `rank_statuses=0,0`. Legacy median was about `5.47 ms` / `46.4k processed tokens/s` (`23.2k unique tokens/s`), while sharded A2A was about `6.72 ms` / `37.8k processed and unique tokens/s`. This is a synthetic native baseline, not Megatron-LM parity.
+- H20 `123.57.26.97:28004`: fresh ABI1 dynamic sharded native smoke passed on both ranks (`rank_statuses=0,0`) against a full-expert reference with complementary source masks. Dynamic grouped-expert parameter/m/v maxima were `1.53e-5` / `4.88e-5` / `5.75e-8`; it also exercised a third tenant with zero global target tokens, clocks `[2,2,0]`, and explicit rejection of ordinary `train_step`. The server path still broadcasts replicated source batches. Dynamic MTP is now covered separately by the restricted single-axis oracle described above.
+- H20 `123.57.26.97:28004`: fresh ABI13 dense base-MLP TP2 smoke passed on both ranks. CLI and server use one CPU sharding helper for gate/up rows and matching down columns; the per-context native flag validates local shapes without process-global env state. C++ reduces the row-parallel output and all-reduces column-parallel input dgrad. Eval/train loss differed from the replicated full-weight reference by `1.84e-5`; FP32 LoRA gradient-accumulator maxima were `3.05e-5` / `4.58e-5`, and the largest post-Adam LoRA slice difference was `9.31e-9`. The smoke also rejects MLP LoRA targets, prevents incompatible TP state transitions, and covers eval-to-train cache invalidation.
+- ABI14 working tree: frozen full attention shards Q/K/V output heads and O input columns, while fixed and selected dynamic Q/K/V/O LoRA use projection-aware layouts. Replicated-side gradients are summed once at the optimizer boundary. Latent-rank LoRA now applies copy-to-TP-region before its local A/B path so a later sharded branch sums input dgrad before it reaches preceding replicated layers. Checkpoint v4 records the layout, replicated tensor geometry, and exact fixed native slot identity; v3 remains loadable only for latent-rank tensors and attention migration is explicitly rejected.
+- H20 target: ABI14 4Q/2KV-head full-attention TP2 smoke passed Q/K/V/O fixed and selected-dynamic full-reference oracles on both ranks. Fixed eval/loss differed by `4.40e-4`; Q/K/V-B and O-A gradient maxima were `1.83e-4` / `1.14e-5` / `3.05e-4` / `1.83e-4`, FP32 m/v maxima were `1.53e-5` / `6.63e-9`, and the local standard-Adam formula error was below `3.73e-9`. Selected-dynamic loss differed by `6.82e-5` and the largest Q/K/V/O parameter difference was `3.05e-5`.
+- H20 target: ABI14 two-layer GDN latent-rank TP2 full-reference smoke passed on both ranks. Eval/loss differed by `5.46e-5`; A/B gradient maxima were `9.54e-7` / `4.77e-7`, FP32 m/v maxima were `4.77e-8` / `3.71e-14`, and the local standard-Adam formula error was zero. ABI14 general single-GPU native smoke and dense-MLP TP2 regression also passed.
+- `99aecef`: ABI16 introduced explicit orthogonal TP/DP groups, dense TP x DP server batch sharding, distributed adapter-registry consensus, and a world4 native TPDP oracle. The H20 `TP2 x DP2` smoke passed fixed and dynamic LoRA with unequal DP token counts and GDN TP2 remained green.
+- ABI17 working tree: CLI and server shard Qwen3.5 tied embeddings and Qwen3.6 untied embedding/LM-head rows on CPU before device transfer. C++ performs masked local embedding plus TP SUM and a two-pass distributed cross entropy with TP MAX/SUM statistics and TP-summed frozen-head dgrad. MTP remains explicitly rejected because its weights are not sharded.
+- H20 target: ABI17 TP2 GDN smoke passed tied fixed-LoRA and untied dynamic multi-LoRA against full-vocabulary references. Fixed loss differed by `5.90e-4`, dynamic loss by `1.16e-3`, FP32 m/v maxima were `1.07e-5` / `1.22e-9`, and the standard Adam formula error was zero. The world4 `dp-tp` rank-order oracle used TP groups `[0,2]` and `[1,3]`, proving vocabulary ranges follow explicit `tp_rank`; local loss differed by at most `2.31e-3`, with FP32 m/v maxima `3.44e-5` / `1.96e-8`.
+- H20 target: ABI17 real-vocabulary GDN benchmark (`batch=1`, `seq=128`, `hidden=2048`, `layers=3`, `vocab=248320`, warmup `2`, iterations `10`) recorded single-rank p50 `24.92 ms`, `5.10k model tokens/s`, and `3.33 GiB` peak resident memory. Vocab-parallel TP2 recorded p50 `23.17 ms`, `5.48k model tokens/s`, and `2.80 GiB` per-rank peak resident memory. The previous replicated-vocabulary ABI16 TP2 baseline was about `24.25 ms` and `4.22 GiB` per rank. This is a native scaling benchmark, not matched Megatron-LM parity.
+- H20 target: GDN backward recurrence fusion is now enabled by default. The fused CUDA template folds the reverse `S_t` undo into the direct-output sweep and removes redundant per-token barriers; `QWEN36_GDN_RECURRENT_FUSION=0` retains the legacy kernel. On the matched synthetic workload (`B=2,S=512,H=2048,L=3`, warmup `5`, iterations `30`), single-rank p50 improved `80.272 -> 75.685 ms` and TP2 p50 `75.908 -> 70.885 ms`; at `B=8`, TP2 p50 improved `157.986 -> 148.216 ms`. Loss and peak resident memory were unchanged. An independent ATen recurrence-backward TP2 smoke retained `adam_error=0` and the same fixed/dynamic loss, m/v, and parameter deltas on both ranks.
+- H20 target: an opt-in two-stage chunkwise GDN backward first computes exact `dS` chunk boundaries, then replays independent `(batch-head, chunk)` gradient CTAs. With `QWEN36_GDN_STATE_CHECKPOINT_STRIDE=32` and `QWEN36_GDN_CHUNKWISE_BWD=1`, TP2 `B=2,S=512,H=2048,L=3` p50 improved `71.484 -> 53.458 ms` (`25.2%`), with repeated `53.432 ms` evidence and unchanged loss; TP2 stride-2 fixed/dynamic oracles retained `adam_error=0`. The opt-in flag and checkpoint stride are part of the distributed topology hash, invalid checkpoint settings fail fast, and checkpoint offsets use `size_t`. It remains opt-in because single-rank p50 regressed `76.195 -> 81.843 ms`, TP2 `B=8` improved only `148.992 -> 147.277 ms`, and state checkpoints add about `100 MiB/GPU` for this workload.
+- The EP HTTP server now has an opt-in bounded cross-request multi-LoRA coalescer. Requests must explicitly set `allow_aggregate_loss=true`, use the same session/sequence/source layout and LoRA rank, and have disjoint adapter IDs; source-major rows are rebuilt into the existing heterogeneous native batch. The default remains one request per native loss so existing clients keep request-local loss semantics. ABI25 requires the native report symbols and returns global per-adapter losses by reducing token-loss numerators/counts only over DP and source-sharded EP; TP and replicated EP are not double-counted. Older native libraries are rejected at load time, all ABI25 ranks negotiate report versus legacy calls before report-only collectives, and a mixed-call TP2 x EP2 smoke fails closed before a subsequent report step succeeds. Coalesced responses retain the explicitly labelled aggregate scalar and expose the requesting adapter range through `adapter_losses`. IPC slab wire version 2 rejects old parent/worker layouts; the server caps a coalesced batch at 2048 adapters for the fixed 256 KiB result slot, oversized worker results become a compact signalled error, and the coordinator rejects rank-inconsistent report vectors. Local tests cover row ordering, shared-row expansion, result slicing, duplicate/rank rejection, opt-out isolation, FIFO dispatch, and IPC overflow/consistency handling; H20 TP2 fixed/dynamic report smoke retained `adam_error=0`.
+- Multi-LoRA HTTP capability v1 now provides an explicit versioned negotiation path: `X-Rustrain-Multi-LoRA-Capability: v1` enables the same bounded coalescing contract as `allow_aggregate_loss=true`, and responses advertise `capability_version=1` plus per-adapter loss, optimizer-step, and coalesced-scope capabilities. Legacy clients retain request-local dispatch semantics. The parser and full local server test suite (50 tests) pass.
+- The EP server now exposes `/v1/sessions/{id}/train_multi_binary`, a version-1 `RLM1` little-endian wire with exact tensor byte lengths, adapter IDs, and optional expected optimizer steps. It dispatches the existing validated IPC slab without JSON/base64 decode; parser rejection tests and the full local server suite (52 tests) pass, and the H20 ABI1 server suite also passes 52/52. Native host-i64 now has opt-in pinned staging, but it did not improve matched H20 p50 and cross-request H2D/compute overlap remains open.
+- Multi-LoRA coalescing now enforces `RUSTRAIN_MULTI_LORA_BATCH_MAX_RANK_WORK` (default `32768`, hard cap `4194304`) on the padded `adapter_count * max_lora_rank` work estimate. Rank-heterogeneous requests can still share a window, but rank inflation seals/rejects the window before dispatch rather than causing an unbounded padded GEMM. The local capacity/overlap tests remain green.
+- Multi-LoRA batching deadlines now start when the first request is admitted, not when its FIFO dispatch job reaches the head of the queue. Time spent running the previous native GPU step therefore consumes the next window's default 2 ms coalescing delay; an already-expired window dispatches immediately instead of creating a deterministic post-compute bubble. This preserves the single-consumer IPC/NCCL ordering and does not claim concurrent GPU steps.
+- Requests carrying `expected_steps` now remain exclusive windows. Their optimistic-concurrency preflight is a request-level failure boundary, so a stale retry no longer fails unrelated tenants flattened into the same native command; clients without step guards retain the bounded coalescing path. Unknown adapter IDs without `expected_steps` still require a future per-request collective validation bitmap.
+- Variable-split EP A2A now has a compact metadata transfer path: `QWEN36_EP_A2A_COMPACT_COUNTS=1` copies only the local send row, receive column, and validity flags (`O(world)` host bytes) instead of the complete count matrix; it defaults on for communicator worlds >=8 and remains opt-in for smaller worlds. H20 TP2 x EP2 x DP2 world8 `tri-smoke` passed with the default compact path. Forced EP2 A/B was `22.707 ms` p50 compact versus `22.385 ms` legacy, so the small-world default remains legacy; true GPU-only variable-count enqueue still requires a DeepEP-like prebuilt dispatcher.
+- ABI18 working tree: native topology owns independent TP, EP, and expert-DP communicators and validates `WORLD_SIZE=TP_SIZE*EP_SIZE*DP_SIZE`. Dense LoRA parameters synchronize and reduce over EP and expert-DP replicas; routed expert parameters remain local to EP and reduce only over expert-DP. TP-latent expert LoRA copies its input into the TP autograd region so backward sums the input gradient exactly once. CLI weight loading composes attention/GDN/vocabulary TP shards with expert EP shards; the server explicitly rejects the source-sharding contract it cannot yet represent.
+- H20 target: ABI18 TP2 x EP2 sharded-A2A smoke passed fixed Q and grouped expert gate-up/down LoRA against a full-model reference with unequal EP token counts. Standard Adam parameter oracles were zero on all four ranks; fixed FP32 m/v maxima were `1.01e-4` / `3.08e-9`. Selected dynamic multi-LoRA produced positive selected updates, exactly zero unselected updates, and clocks `[1,0]`; invalid topology and MTP guards fired.
+- H20 target: ABI18 regressions passed TP2 x DP2, GDN TP2 fixed/dynamic, and pure EP2 sharded A2A fixed/dynamic tests. GDN fixed/dynamic FP32 m/v maxima remained within `1.07e-5` / `1.22e-9`, and all standard Adam formula errors were zero. Pure EP fixed expert parameter differences were bounded by BF16 (`9.62e-4` maximum) with m/v maxima `1.22e-5` / `3.92e-9`.
+- ABI19 packs all top-k assignments into one sharded EP dispatch/combine, sorts received rows once, and schedules base and aligned fixed-LoRA expert projections with prebuilt PyTorch `_grouped_mm`. Base grouped GEMM no longer falls back when the local LoRA rank is unaligned; only that low-rank delta uses per-expert matmul, followed by one packed TP reduction. `QWEN36_EP_A2A_PACKED=0` retains the old routing-slot loop for diagnosis.
+- H20 target: ABI19 default-packed TP2 x EP2 fixed/dynamic smoke passed. The unaligned local-rank branch retained the ABI18 full-reference bounds (`1.98e-3` maximum BF16 parameter difference, `9.92e-5` / `3.02e-9` FP32 m/v maxima, zero standard-Adam error), and dynamic adapters kept positive selected updates, exactly zero isolated updates, and clocks `[1,0]`.
+- H20 target: matched native TP2 x EP2 fixed-LoRA benchmark (`B=2`, `S=128`, `H=1024`, `E=8`, `I=2048`, top-k 2, global rank 16, one layer, warmup 5, iterations 30) aggregates each measured step with a world-wide maximum across the TP x EP grid. It reduced p50 step time from `9.007 ms` to `6.499 ms` and increased unique-token throughput from `56.84k/s` to `78.78k/s` (`27.8%` lower p50, `38.6%` higher throughput). Allocator peak stayed near `0.26 GiB`; the larger packed workspace increased observed resident memory by at most `0.09 GiB` in this run.
+- H20 target: matched Nsight Systems traces recorded `6812 -> 5168` `cudaLaunchKernel` calls, NCCL SendRecv `96 -> 48`, AllGather `24 -> 12`, and BF16 AllReduce `216 -> 120` across four ranks for one warmup plus two measured steps. These traces predate and therefore exclude the benchmark-only world-max metrics collective, which runs outside step timing. The counts include initialization, so they are conservative end-to-end process-tree evidence rather than isolated dispatcher counts.
+- ABI20 shards routed and shared expert gate/up output rows and down-projection input columns across TP ranks. Projection LoRA follows the same column/row layouts, reducing only replicated A or B gradients at the optimizer boundary. Packed EP dispatch computes local ETP partials, TP-reduces them before routing weights, and no longer duplicates frozen expert compute or storage across TP peers. Checkpoint v5 now encodes the routed expert's compound EP leading axis plus segmented gate/up ETP axis; the server permits this topology only with packed sharded A2A and rejects replicated-expert TP early.
+- H20 target: ABI20 TP2 x EP2 fixed and dynamic smokes passed routed gate/up/down, shared gate/up/down, selected-tenant isolation, and FP32 Adam-state oracles. Pure EP2, TP2 x DP2, and GDN TP2 regressions also passed. On the matched ABI19 workload, replicated experts recorded p50/p95 `8.742/9.822 ms` and `58.57k` unique tokens/s; ETP recorded `6.923/7.874 ms` and `73.95k` unique tokens/s, a `20.8%` p50 reduction and `26.3%` throughput increase. Peak allocator allocation fell from about `0.191 GiB` to `0.153 GiB`, while the frozen expert replication factor fell from two to one.
+- The CLI and server now share checkpoint v5 and the deterministic standard-PEFT merger. Distributed CLI export compacts active native slots, preserves projection-aware layouts and `[A,B]` optimizer ordering, writes isolated rank shards under a unique attempt, and atomically publishes a complete PEFT directory. The fixed-LoRA CLI accepts `--resume-from` for the same topology and restores LoRA tensors, FP32 Adam m/v, loss, and logical step after validating the base model and adapter signature.
+- Checkpoint v5 replaces the TP-only v4 contract with explicit rank order, five-dimensional coordinates, multi-axis TP/EP placements, fused routed gate/up segments, and fixed/dynamic slot identities. Every distributed rank writes an isolated directory; v3/v4 remain read-only compatible for their original TP subset. A unique staging-directory fence, rank error files, a completion marker, and directory rename prevent attempt reuse or partial artifact publication. EP IPC dispatch is serialized and propagates errors from every worker.
+- H20 target: ABI20 TP2 x EP2 fixed/dynamic native smoke remained green after v5. The v5 TP2 x EP2 sentinel completed save, per-rank fresh restore, and fused expert merge with global `[gate_all|up_all]` ordering; the dynamic merge preserved tenant identity and optimizer clock. Local server tests and the H20 ABI1 run both passed the merge cases.
+- H20 target: fixed TP2 x EP2 state was copied to CPU to simulate safetensors, restored into a fresh context, and continued for one step. All four ranks reported zero loss-path parameter/Adam m/v divergence from uninterrupted training and optimizer step `2`; dynamic selected-tenant isolation remained green.
+- ABI20 working tree: native Qwen sessions compose TP, EP, and expert-DP coordinates into one process grid and use the EP x DP source coordinate only for source-sharded A2A. Checkpoint v5 now binds atomically published manifests to tensor-file content digests and an explicit per-save generation supplied by the distributed coordinator. New writers publish a constant-size rank receipt with manifest/schema/shared-metadata/file digests; each loader preflights `O(world_size)` receipts and parses/hashes only its full local manifest/state, while digest-v5 checkpoints without receipts retain the full-manifest compatibility fallback. Missing/mixed/reused generations, missing receipts, and divergent DP adapter or optimizer replicas are rejected. The fixed Adam clock is recorded separately from the session transport step.
+- H20 target: TP2 x EP2 x expert-DP2 world8 fixed/dynamic smoke passed all ranks. Fixed parameters matched the full-rank oracle within `1.96e-3`, FP32 Adam m/v within `1.33e-6` / `1.79e-11`, and the standard Adam formula exactly. A CPU-state fresh resume continued with zero parameter/m/v divergence and fixed step `2`. Selected dynamic training kept the unselected tenant unchanged, preserved fixed-clock isolation, used tenant clocks `[1,0]`, and passed the full-rank parameter/Adam oracle plus exact second-step resume. The world4 TP2 x EP2 regression remained green.
+- H20 target: TP2 x EP2 x expert-DP2 replicated-source world8 also passed. It retained attention/vocabulary TP, used latent-rank MLP LoRA with replicated expert base weights, and exercised the distinct replicated A2A normalization path. Fixed Adam error was zero, fixed parameter/m/v differences were at most `1.96e-3` / `2.58e-6` / `2.15e-11`, and fresh resume was exact at step `2`. Dynamic tenant isolation and Adam error were zero; full-reference parameter/m/v differences were at most `1.99e-3` / `4.16e-5` / `5.13e-10`.
+- Server source-sharding working tree: EP workers derive source rows from the configured five-axis topology. TP peers at the same `(DP,EP)` coordinate receive identical rows; source-sharded A2A uses `dp_rank * ep_size + ep_rank`, while replicated-source EP uses only `dp_rank`. Multi-LoRA source-parallel requests require the complete `n_total * source_parallel_size` global batch instead of silently duplicating `n_total` rows across DP ranks. Server weight loading now composes EP expert narrowing with routed/shared expert TP before vocabulary and attention/GDN sharding, matching the CLI path.
+- Server IPC working tree: a command has one absolute deadline; an incomplete semaphore generation permanently poisons the channel, the first IPC failure terminates and reaps the owned worker group, later dispatches fail immediately, and EP health changes to unavailable. Partial worker launch failures also terminate and reap ranks already started. A singleton worker-group session contract now rejects a second session ID and actually clears rank-local state on delete; tenants within the base model continue to use independent dynamic LoRA adapters. EP train responses return the real logical step instead of `0`.
+- H20 target: the fresh source-sharded world8 native oracle remained green after the server integration. Fixed resume restored parameters and Adam m/v with exact zero divergence at step `2`; selected dynamic training reported zero isolated-tenant change, zero standard-Adam error, and clocks `[1,0]` on every rank. Rust server/IPC tests for that milestone ran in the local ABI1 host environment; the target's later rustup toolchain was not used for those historical results.
+- ABI21 working tree: distributed server checkpoint save/load now holds one coordinator dispatch lock across all phases. Save writes to an exclusive sibling staging directory, validates the complete receipt set by hydrating rank-local shadow contexts, aborts those shadows, and publishes only a fresh destination with `RENAME_NOREPLACE`. Load builds communicator-attached shadows without broadcasting temporary random parameters, hydrates fixed/dynamic LoRA and Adam clocks in isolation, and swaps live contexts only after every rank prepares successfully; an incomplete commit fail-stops the worker group. Existing destinations, immutable generation/`LATEST` rollover, power-loss fsync durability, and cross-topology reshard remain open.
+- H20 target: ABI21 source-sharded TP2 x EP2 x expert-DP2 world8 shadow resume passed all ranks. DP-nonzero fixed tensors deliberately diverged before no-sync communicator attach and remained bitwise unchanged after attach; checkpoint hydrate then restored fixed/dynamic A/B, FP32 Adam m/v, and tenant step exactly. Continued fixed and selected-dynamic training retained zero resume divergence. The native library exported the ABI version, no-sync attach, and restore-adapter symbols from the prebuilt ABI1 PyTorch/NCCL build.
+- ABI23 adds live heterogeneous adapter registration and a selected v2 trainer that deterministically groups adapters by rank, canonical active slots, and local tensor geometry. Each group reuses the homogeneous training path; if a later group fails, persistent out-of-place Adam shadows restore every earlier group's parameters, m/v, and tenant clock. The Rust tensor binding and production borrowed host-i64 server path both call v2 while the existing HTTP/IPC rank field remains wire-compatible. H20 single-rank smoke passed success, legacy rejection, duplicate/unknown-ID registry preservation, host-i64 dispatch, and injected cross-group rollback; TP2 x EP2 x expert-DP2 world8 passed rank-3 `q_proj` plus rank-4 multi-module heterogeneous selection on all ranks, including rank-0-only wrapper and pre-commit Adam failures that all ranks observed before the next group. Per-signature forward/backward, synchronization, Adam launches, and GPU-to-host success fences remain a performance gap until a unified finalizer is implemented. The consensus boundary prevents ranks whose homogeneous trainer returned from entering different subsequent groups; it is not recovery from an unrecoverable rank-local CUDA/NCCL failure inside that trainer.
+- ABI24 adds packed LoRA gradient synchronization. `GroupedLoraGradientSyncPlan` groups contiguous FP32 accumulators by EP/DP/TP reduction mask and issues one NCCL all-reduce per populated bucket/axis; `QWEN36_PACKED_LORA_SYNC=0` retains the per-tensor grouped fallback, and the setting is part of the distributed topology hash. H20 TP2 fixed-LoRA smoke and TP2 x DP2 oracle passed with packed mode enabled and disabled. On the synthetic `B=2,S=512,H=2048,L=3,rank=8` TP2 benchmark, p50 was `77.970 ms` packed versus `78.024 ms` fallback (noise-level improvement), so this reduces collective launch count without claiming backward overlap; token-count CPU fences and full-backward synchronization remain open.
+- ABI24 GDN backward can optionally checkpoint exact FP32 recurrent state at fixed token boundaries (`QWEN36_GDN_STATE_CHECKPOINT_STRIDE`, default `0`). Reverse chunks restart from their exact end state, limiting reconstruction error accumulation across chunks while preserving strict-right-padding lengths. H20 TP2 smoke passed with `stride=2` across the `S=9` full-weight gradient/Adam oracle. On the matched `B=2,S=512,H=2048,L=3,rank=8` TP2 benchmark, `stride=128` changed p50 from `76.020 ms` to `76.151 ms` (`+0.17%`) and peak observed resident memory from `1.8086 GiB` to `1.8477 GiB` (`+40 MiB/GPU`). The path remains opt-in because it stores `ceil(S/stride)+1` states per layer; it still inverts clamped decay inside each chunk and is not sequence-parallel GDN.
+- Server batch registration now rejects non-positive or over-limit `BatchAddLora.count` in both HTTP and worker paths (`RUSTRAIN_MAX_BATCH_ADD_LORA`, default 64, hard cap 4096) before allocation. If a later adapter allocation fails, each worker rolls back its successfully created IDs in reverse order, preventing partial registry state from surviving a failed batch. This remains a safety/consistency fix rather than tenant batching; the scheduler described below still serializes worker-group collectives.
+- Server EP dispatch now uses a bounded FIFO single-consumer scheduler for control-plane and checkpoint jobs (`RUSTRAIN_EP_DISPATCH_QUEUE_CAPACITY`, default 32, hard maximum 4096). Each accepted job runs on one `spawn_blocking` worker at a time, preserving the existing IPC lock, timeout poison, and checkpoint transaction boundaries; caller disconnect does not cancel accepted work. Queue pressure returns HTTP `429`, scheduler/worker failure returns `503`. Tensor train/eval routes still admit only one 48 MiB request and return `429` while busy, so this is runtime backpressure/fairness rather than cross-tenant GPU batching. Local server tests covered FIFO, single in-flight, queue-full, current-thread runtime liveness, panic recovery, caller disconnect, and status mapping; they were not rerun on the H20 target for that milestone.
+- Heterogeneous selected v2 training now defaults to one padded activation batch (`QWEN36_HETERO_PADDED_BATCH=1`). Each projection pads only its LoRA rank dimension to that projection's selected maximum, represents inactive target holes with zero tensors, preserves each tenant's `alpha / logical_rank` scaling, and runs one TrainOnly call followed by one FinalizeOnly call. `QWEN36_HETERO_PADDED_BATCH=0` retains deterministic signature grouping and cross-group rollback. H20 single-GPU and TP2 x EP2 smokes passed rank-8/rank-9 and target-hole updates in padded mode; the four-rank grouped fallback also passed. On a matched TP2 x EP2 synthetic heterogeneous workload (`B=4,S=128,H=1024,E=8,I=2048`, ranks 16/8, four tenants, warmup 3, iterations 10), padded mode changed p50 from `17.750 ms` to `13.841 ms` and unique-token throughput from `57.69k/s` to `73.98k/s` (`22.0%` lower p50, `28.3%` higher throughput). Peak allocator allocation rose from about `0.240 GiB` to `0.354 GiB`; p95 was `20.152 ms` padded versus `19.441 ms` grouped in this short run, so tail-latency improvement is not established.
+- Dynamic multi-LoRA train requests now optionally carry one `expected_steps` value per selected adapter. The native validation hashes ordered IDs and expected clocks across TP/EP/DP before checking local clocks, so stale retries fail collectively before any optimizer update. Coalesced HTTP windows reject mixing guarded and unguarded requests, and responses return the actual per-adapter optimizer steps so clients can continue safely even when a tenant receives zero valid tokens. Legacy commands and native results deserialize with empty defaults. H20 ABI1 single-rank smoke rejected a stale clock while leaving both tenant clocks at zero, and TP2 x EP2 smoke rejected a rank-0-only expected-clock divergence on all four ranks before completing the existing heterogeneous training oracle.
+- H20 ABI1 Qwen3.6 MTP extension: the TP2 x EP2 sharded-MoE dynamic oracle now covers uneven source rows (global main/MTP tokens `7/5`) and matches the full-source fixed-MTP oracle with loss diff `4.38e-4`, parameter diff `2.01e-3`, and isolated-tenant state diff `0`. A TP2 vocabulary-parallel MTP smoke also passed with loss diff `3.62e-3`, Adam `m/v` diffs below `1.8e-4`/`1.1e-8`, and zero unselected-tenant change. The report path keeps MTP token sums under the main-loss denominator while gradient normalization uses the global main/MTP ratio; PP/CP MTP remains fail-closed.
+- H20 ABI1 selected-v2 registry optimization: restore now puts selected adapters back by saved canonical index, and a persistent sorted `adapter_id -> canonical slot` index serves the distributed validator, selected-v2, expected-step, and selected-eval paths without rebuilding or scanning the full registry each step. Temporary selected/chunk registries invalidate the index and restore it only with the canonical vector. On TP2 x EP2 synthetic dynamic-MoE runs (`B=8,S=128,H=1024,I=2048`, rank 16, q/expert targets, 50 iterations), 8 registered/8 active tenants measured `13.859/17.561 ms` p50/p95; 1024 registered/8 active measured `14.390/19.579 ms` (about `3.8%` p50 overhead) while allocator reserved memory increased from about `1.04` to `9.46 GiB/GPU`. Local smoke and TP2 x EP2 distributed smoke passed. This is a registry-scaling result, not a complete-model throughput claim; the memory cost remains an open capacity constraint.
+
+# Megatron-LM Performance Gap
+
+ABI20 removes frozen routed/shared expert replication across TP ranks, but it
+does not close the Megatron MoE throughput gap. Rustrain's replicated TP input
+contract uses column-parallel FC1, row-parallel FC2, and a TP SUM before routing
+weights; Megatron's sequence-parallel dispatcher instead composes EP A2A with
+expert-TP gather/reduce-scatter. Split planning still has host-visible counts,
+dispatch and grouped GEMM use the current stream without communication overlap,
+and permutation/combine remain separate ATen operations. Capacity-aligned
+routing, fused permutation and inverse permutation, shared-expert overlap, a
+DeepEP-class backend, sequence parallelism, cross-topology checkpoint resharding,
+and a matched Megatron benchmark remain future work. The current evidence is a
+matched replicated-expert-versus-ETP native result, not Megatron-LM parity.
+
+# Decisions During Execution
+
+- Keep TP and EP communicators separate; do not reuse the existing EP `LayerConfig.nccl_comm` for LoRA TP deltas.
+- Publish multi-LoRA `n_max` directly from rank 0 with `ncclBroadcast`; filesystem rendezvous can reuse stale files across process restarts and give ranks different chunk schedules.
+- Do not enable Qwen TP/PP/CP by merely relaxing runtime validation.
+- Treat dense and expert MLP TP as projection-aware layouts: gate/up replicate A and shard B, while down shards A and replicates B. Reduce only the replicated parameter gradient at the optimizer boundary; never apply a replicated-projection reduction rule to a disjoint shard.
+- For frozen full-attention TP, use Q/K/V output-head sharding and O input-column sharding. Q/K/V replicate LoRA A and shard B; O shards A and replicates B. Sum only the replicated-side gradient at the optimizer boundary so the activation collective is not duplicated.
+- Do not reinterpret v3 attention LoRA tensors as projection-aware shards. Their latent-rank geometry is incompatible, so require v4 for Q/K/V/O resume and keep v3 compatibility for latent-rank-only modules.
+- Treat Exa/Jina dependency search failures as missing evidence, not as proof that a package is compatible.
+
+# Verification
+
+Passed: `cargo check -p rustrain-core -p rustrain-qwen3-6 -p rustrain-server` (with the repository host venv), focused runtime/topology unit tests, remote ABI8-ABI14 native smokes listed above, ABI16 TP2 x DP2 and GDN TP2 smokes, ABI17 TP2 tied/untied vocabulary parity plus world4 custom-rank-order TPDP smokes, ABI18 TP2 x EP2, TP2 x DP2, GDN TP2, and pure EP2 sharded-A2A native smokes, ABI19 default-packed TP2 x EP2 fixed/dynamic smoke plus fixed/dynamic benchmarks and matched profiler traces, and ABI20 ETP TP2 x EP2 fixed/dynamic smoke and matched replicated-versus-sharded expert benchmark. ABI20 regressions passed TP2 x DP2, GDN TP2, pure EP2 fixed/dynamic, source-sharded and replicated-source TP2 x EP2 x expert-DP2 world8, and CPU-checkpoint fresh-context resume smokes. Checkpoint v5 local/H20 tests passed pure EP rank isolation, custom rank order, TP2 x EP2 save/fresh-restore/fused-expert merge, TP2 x EP2 x expert-DP2 compact-receipt preflight and replica validation, stale-manifest content rejection, independent fixed/dynamic optimizer clocks, dynamic tenant identity/clock merge, CLI shared export/resume contracts, digest-v5 fallback, and v3/v4 read compatibility. Heterogeneous padded single-GPU and TP2 x EP2 smokes, the grouped TP2 x EP2 fallback, the matched padded/grouped benchmark, the ABI1 router-aux checkpoint smoke, ABI29 dense GDN CP2 fused-disabled/fused-enabled full-reference smoke and SendRecv trace, ABI29 `CP2 x PP2` dynamic lifecycle mismatch/recovery, and ABI29 PP2 active-window mutation rejection passed with the prebuilt H20 PyTorch/NCCL runtime.
+
+Not run: interleaved/chunked PP, hybrid GDN CP>2 or multi-axis CP, MTP combined with CP/PP, real multi-node training, end-to-end Rust HTTP checkpoint transactions on the H20 target, concurrent HTTP request coalescing into a native heterogeneous batch, cross-topology resharding, and a matched Megatron performance benchmark. Dense full-attention ring CP2/CP4 correctness has been run, but its serialized ATen implementation has no throughput claim. The fixed/dynamic-LoRA H20 PP2, fixed-LoRA PP3, dense GDN CP2, restricted dynamic-MTP, and CP2/CP4 ring smokes used the existing prebuilt PyTorch ABI1 and NCCL runtime; no dependency installation or JIT workaround was used. Rust IPC/server tests remain local-only evidence; the current target rustup toolchain has only been used for the focused ABI1 Qwen build/test path.
diff --git a/docs/plans/qwen-lora-megatron-spec.md b/docs/plans/qwen-lora-megatron-spec.md
new file mode 100644
index 00000000..bcb3e1c8
--- /dev/null
+++ b/docs/plans/qwen-lora-megatron-spec.md
@@ -0,0 +1,63 @@
+---
+type: ChangeSpec
+title: Qwen3.5/3.6 LoRA Megatron-level runtime
+description: Contract for LoRA-only distributed parallelism, MoE performance, GDN, and independent multi-tenant training.
+tags: [qwen3.5, qwen3.6, lora, megatron, distributed]
+timestamp: 2026-07-17T00:00:00Z
+---
+
+# Problem
+
+The native Qwen3.5/3.6 path has correct and tested single-rank, EP, and dense-DP slices, but it is not yet a Megatron-LM-level LoRA runtime. In particular, Qwen native TP/PP/CP are not implemented, and the MoE path lacks fused asynchronous token dispatch.
+
+# Target Outcome
+
+Provide a Qwen3.5/3.6 LoRA-only training backend whose parallel groups, rank-local adapter parameters, optimizer/checkpoint state, and runtime scheduling are correct for the supported TP/PP/DP/EP/CP topology. Base-model parameters remain frozen, but every trainable LoRA projection and every tenant update must be independent and reproducible.
+
+# Contract
+
+- Rust owns configuration, scheduling, data partitioning, checkpoint orchestration, and transport. C++/CUDA owns all training-path math and collectives.
+- No Python JIT or ad-hoc dependency builds. A dependency is used only after a compatible prebuilt package is found and verified on the target runtime.
+- TP and EP communicators are distinct. EP routed-output collectives must never be reused for LoRA TP output reduction.
+- A global LoRA rank is partitioned over TP ranks (`rank % tp_size == 0`); rank-local A/B shapes and checkpoint metadata identify the topology.
+- PP owns disjoint layer ranges and a real microbatch schedule. CP owns sequence/KV exchange; configuration fields alone do not satisfy either contract.
+- Each dynamic adapter has independent optimizer step, m/v, gradient accumulation state, and checkpoint metadata. A request may train a selected subset of tenants without mutating others.
+
+# Requirements
+
+1. Correct Qwen3.5 full-attention and Qwen3.6 hybrid GDN/full-attention forward/backward with LoRA target projections.
+2. Correct LoRA-only TP, DP, EP, PP, and CP group semantics, with explicit unsupported combinations rejected before model allocation.
+3. MoE grouped/fused dispatch and communication overlap where a verified prebuilt dependency exists; otherwise retain a correct native fallback and report the gap.
+4. FP32 gradient accumulation and logical optimizer boundaries, including abort/zero behavior.
+5. Dynamic multi-LoRA selected-tenant scheduling, independent optimizer clocks, and no cross-tenant parameter or state updates.
+6. Rank-aware checkpoint save/load and resume equivalence.
+7. Target-machine smoke, numerical parity, throughput, peak-memory, and scaling evidence against a Megatron-LM reference under matched settings.
+
+# Out Of Scope
+
+Full-parameter training, unfrozen base weights, unverified third-party packages, Python JIT extensions, and claiming Megatron parity from topology metadata or unit tests alone.
+
+# Acceptance Evidence
+
+| Criterion | Direct evidence |
+| --- | --- |
+| TP LoRA rank sharding | two-rank native smoke checks local A/B shapes, TP delta all-reduce, and finite loss/update |
+| DP/EP separation | multi-rank smoke checks replicated LoRA reduction and local expert ownership separately |
+| PP/CP | multi-rank stage/ring smoke with layer ownership and sequence/KV exchange assertions |
+| Tenant independence | selected-adapter test proves untouched tenant tensors, m/v, and step are unchanged |
+| Accumulation | FP32 accumulation test matches concatenated-batch reference and abort clears pending state |
+| Checkpoint resume | save/load continuation matches uninterrupted optimizer state and bias correction |
+| Performance | target-machine matched benchmark records tokens/s, step time, peak memory, and communication share |
+
+# Dependency Constraints
+
+The target H20 runtime exposes PyTorch 2.12.1+cu130 with C++ ABI1,
+CUDA 13.0, and NCCL 2.29.7. Transformer Engine, flash-attn, FLA, DeepEP,
+and grouped-gemm packages are not installed. Triton and Tilelang import, but
+their runtime compilation model is excluded by the no-JIT contract. The
+current implementation therefore uses the prebuilt libtorch `at::_grouped_mm`
+and NCCL APIs without adding a dependency.
+
+# Known Baseline
+
+The repository currently has a working native EP/dense-DP subset, grouped MoE microbenchmark evidence, and a separate 5D topology helper. These are prerequisites, not proof of the target outcome.
diff --git a/docs/qwen35-qwen36-megatron-audit.md b/docs/qwen35-qwen36-megatron-audit.md
new file mode 100644
index 00000000..2b412e64
--- /dev/null
+++ b/docs/qwen35-qwen36-megatron-audit.md
@@ -0,0 +1,308 @@
+# Qwen3.5/3.6 LoRA 并行与性能审计
+
+本文记录当前 native Qwen3.5/3.6 LoRA 后端与 Megatron-LM 级训练栈的边界。结论按实际代码和 smoke/integration 结果整理，不把配置字段或通用拓扑类型当作已经实现的 kernel。
+
+本轮补上了 native host-i64 入口的 opt-in `QWEN36_HOST_PINNED_STAGING=1`：context 缓存独立的 pinned input/target/attention staging，并以 non-blocking H2D 替代 pageable blocking copy。H20 TP2 x EP2 matched native benchmark（`B=2,S=64,H=1024,E=8,I=2048,L=1`，5 warmup/50 iterations）loss 保持 `5.411833`，pageable/pinned p50 为 `6.336/6.959 ms`，p95 为 `7.148/7.108 ms`；因此默认关闭，且不能把该路径等同于已经实现 request/compute overlap。
+
+## 结论
+
+- 模型语义：Qwen3.5 dense、Qwen3.6 dense/MoE 的 native forward/backward 路径已经覆盖 hybrid full attention、GDN、MoE、MTP 和 LoRA 目标模块；已有配置解析、合成 oracle、集成测试及 H20 native smoke 证据，但尚未完成真实 35B/3.6 权重的长时间训练验证。
+- 已实现并可验证的分布式子集：ABI26 可按 `TP-CP-EP-DP-PP` 五维拓扑建立、缓存和 attach 正交 NCCL process groups；LoRA 模型执行已覆盖可组合的 TP、MoE EP 和 expert-DP，包含 TP2 x EP2、TP2 x DP2 和 TP2 x EP2 x expert-DP2 native oracle。固定 LoRA PP 已覆盖固定 shape、单 chunk 的 non-interleaved 1F1B，并在 H20 PP2/PP3 通过。dense full-attention 新增 opt-in arbitrary-CP ring attention：FP32 online-softmax、手写 dQ/dK/dV backward 和 KV/gradient owner-return，H20 CP2/CP4 ring 对 CP1 oracle 通过；Qwen3.6 hybrid GDN 的 CP 仍限制 CP2，CP 与其他并行轴仍 fail-closed。DP 动态租户按 adapter token count 加权，sharded A2A 保留 source flattened row 来恢复租户，并按全局租户 token count 归一化。
+- 性能：ABI19 将 top-k assignment 合并为单次 packed dispatch/combine，并对接收 token 做一次 expert sort 和 grouped GEMM。H20 TP2 x EP2 端到端 fixed-LoRA benchmark 按每步 TP x EP 全局最慢 rank 计时，p50 从 `9.007 ms` 降到 `6.499 ms`，unique-token throughput 从 `56.84k/s` 提升到 `78.78k/s`；这是 native legacy/packed 对照，不是 Megatron 对比。
+- 性能：ABI28 在 heterogeneous selected-v2 入口把 health、request、loss-report capability、report buffer 和 accumulation-clear 状态打包成一个 `int32[6]`，每个 EP/DP/TP 轴只做一次 MIN collective。H20 TP2 x EP2、B=8/S=128/8 active tenants 相对旧六次 collective 实现的三次交错 A/B，中位 p50 约从 `13.394 ms` 降到 `13.199 ms`，Nsight Systems 的 `u32` all-reduce 实例从 `1032` 降到 `752`；这是 native synthetic benchmark，不是 Megatron 对比。
+- 性能：EP A2A count all-gather 后现在只构建一次 CPU `int64` send/receive prefix plan，dispatch backward 与 combine 两个方向复用该 plan，删除重复的 CPU count materialization、memcpy 和 prefix scan。H20 TP2 x EP2 dynamic matched A/B 的 mean 为 `13.974/13.977 ms`、p50 `13.793/14.008 ms`、p95 `15.243/14.736 ms`，显存不变；这是确定性的 host bookkeeping 改进，未观察到稳定端到端加速。
+- 性能：receive-side local expert permutation 新增默认关闭的 `QWEN36_MOE_FUSED_LOCAL_PERMUTE=1`。CUDA counting-sort kernel 直接完成 histogram、prefix scan、metadata scatter 和 inverse row permutation，custom autograd 保存 mapping 供 backward；旧的 `sort/index_select/bincount/cumsum/index_add` 路径保持可用。该 flag 参与 distributed runtime hash，不依赖 TE/DeepEP/JIT。H20 TP2 x EP2 local/distributed smoke 的 fixed Adam error 为 `0`、resume 参数/m/v diff 为 `0`、heterogeneous-v2/recovery-poison 均为 `ok`；ABI1 `cargo check -p rustrain-qwen3-6 --lib` 通过。匹配 dynamic B=8/S=128/H=1024/I=2048/rank16、packed A2A 与 fused weighted-unpermute 的 A/B 中，fused/fallback mean `13.342/13.956 ms`、p50 `13.156/13.846 ms`，但 p95 `15.808/14.792 ms` 未稳定改善，因此默认关闭，也不等同于 Megatron/DeepEP 的通信计算 overlap。
+- 性能：packed MoE 新增默认关闭的 `QWEN36_MOE_FUSED_UNPERMUTE=1`。CUDA kernel 直接按 inverse order 做 `[token, top_k]` 加权归并，custom autograd 同时返回 routed activation 与 routing-weight 梯度；仅替换 packed A2A 的 `index_select`/临时加权张量/`index_add`，legacy/non-packed 路径不变。H20 TP2 x EP2 local/distributed smoke、异构 LoRA、恢复与 poison recovery 均通过，固定 Adam error 为 `0`、resume 参数/m/v diff 为 `0`、dynamic isolation diff 为 `0`；同配置 A/B fused/fallback mean `14.063/14.154 ms`、p50 `13.963/14.132 ms`、p95 `15.248/15.223 ms`，显存不变。这个约 `0.6%` 的均值差异属于噪声范围，不能宣称稳定加速，也不能替代 Megatron/DeepEP 的 fused permutation、GPU-only metadata 或通信计算 overlap。
+- 性能：dynamic transactional Adam 新增默认开启的 `QWEN36_DYNAMIC_ADAM_SHADOW_POOL=1`。每个租户的 Adam `m/v` 仍按租户保留，但 out-of-place rollback shadow 只在 active transaction 期间从 context pool lease，并在 selected-v2 的外层 consensus 后归还；空闲 slot 会按新 layout 重建，避免 rank/target layout 历史种类无限扩容。H20 TP2 x EP2、1024 registered/8 active、完整 128-step tenant rotation：旧路径 pooled 前的 allocated/reserved/max-resident 为 `8.439/8.813`, `9.029`, `11.363 GiB`，pool 后为 `5.028/5.447`, `5.637`, `7.971 GiB`；p50/p95 为 `21.179/32.224 -> 16.752/26.057 ms`。第二个完整周期 allocated 仍为 `5.028 GiB`；pool-disabled fallback、local/distributed smoke 和回滚/poison 检查均通过。该优化解决 shadow 随历史活跃租户增长的问题，但持久 Adam state 仍会增长，尚未达到 Megatron 式 CPU/GPU optimizer-state paging。
+- 性能：dynamic multi-LoRA 的 activation chunk 现在在 device 上累积 loss、hidden-gradient finite 和 MTP token-count validity，整个 logical step 末尾只保留一次 host read/collective consensus；selected loss report 的一次 D2H 也移动到 finalizer 前。这样把数值检查从 `O(chunks)` 个 host fence 降为 `O(1)`，同时保持 Adam/clock commit 前的 fail-closed 语义。H20 TP2 x EP2 uneven-source MTP、TP2 fused-MoE MTP、TP2 vocabulary-parallel MTP 和 EP dynamic regression 均通过；这属于 runtime synchronization 优化，不是 backward bucket overlap 或通信计算重叠。
+- 已实现 LoRA latent-rank TP、frozen full-attention/GDN/dense SwiGLU MLP TP、routed/shared ETP，以及 embedding/LM-head/vocabulary TP。GDN 按 K/V head group 切分复合 QKV、depthwise conv、Z/A/B、A_log/dt_bias 和 out-proj input columns；fixed 与 selected dynamic LoRA 使用 projection-aware shard 和梯度 reduction。ABI20 对 routed fused gate/up 的 gate/up 两半分别切分后重排，并在 routing weight 前归约 local expert partial。TP 可与 EP 或 expert-DP 组合；新增 `QWEN36_SEQUENCE_PARALLEL=1` 的 guarded TP2 dense path，fixed 和 projection-aware dynamic LoRA 均可按 Megatron 语义执行 sequence scatter/all-gather/reduce-scatter/loss gather；CP/PP/EP/DP、MoE、MTP、aux loss，以及 latent-rank dynamic target 仍 fail-closed。
+- CLI 与 server 共享 checkpoint v5 实现。分布式 CLI 训练按 shared run ID 隔离 rank 日志，以唯一 save transaction generation 协调 rank shard，原子发布标准 PEFT 目录，并支持相同 topology 下恢复 fixed LoRA、FP32 Adam m/v 和独立 fixed optimizer step。checkpoint library 不再把 run-scoped attempt ID 隐式当作 save generation；缺失、空、非法或复用 generation 的分布式保存直接拒绝。新 writer 用内容 digest 绑定 manifest/tensor，并发布定长 compact rank receipt；每个进程预检 receipt set 后只解析、hash 和加载本地完整 manifest/tensor。distributed server 的 save/load coordinator 在同一个 dispatch lock 内执行两阶段事务：save 写入 exclusive sibling staging，要求所有 rank receipt 完整并用 shadow context 全量 hydrate 验证，再以 `RENAME_NOREPLACE` 发布 fresh destination；load 先在所有 rank 建立 no-sync shadow context，全部成功后才交换 live context，commit 不完整则 worker group fail-stop。当前不覆盖已有 destination，不提供 `LATEST` generation 指针或断电级目录 fsync 保证；跨 topology reshard 仍未实现。
+- ABI22 server 训练数据面在 parent/worker IPC 间使用 64-byte aligned binary tensor slab：parent 只做一次 base64 decode/packing，worker 零拷贝借用 host `int64` span，按 topology 选取本地 rows，并通过一次粗粒度 C++ 调用完成 H2D 与完整 train/eval step。shared-memory layout、epoch、span 边界/重叠和 timeout poison 均有校验；默认 slab 为 32 MiB，可通过环境变量配置。tensor routes 在 body decode 前使用单许可 admission gate，与本来串行的 IPC dispatch 对齐并限制并发请求内存。HTTP 边界仍是 JSON/base64，host memory 也尚未 pin，因此这不是最终的高吞吐 ingress。
+- 因此当前实现不能宣称“Megatron-LM 级别”。它是一个计算集中在 C++ 的 LoRA TP/EP/expert-DP 子集，离 Megatron 的完整并行和通信重叠仍有实质差距。
+
+## 当前能力矩阵
+
+| 能力 | 当前状态 | 证据/限制 |
+| --- | --- | --- |
+| Qwen3.5 full attention | 已实现 | native smoke、Qwen3.5 配置集成测试 |
+| Qwen3.5 GDN/linear attention | 已实现 | CUDA delta-rule forward/backward 与 native smoke |
+| Qwen3.6 MoE | 已实现 | grouped dispatch、EP smoke；完整模型仍需目标 GPU/权重运行 |
+| MTP | 已实现受限子集 | fixed LoRA 和 dynamic selected LoRA 的 C++ hidden gradient；单轴 DP、TP2 dense/真实 MoE prediction-layer shard、TP2xEP2 不均匀 source-token oracle 均有 H20 证据。TP2 MTP 现在支持 vocabulary-parallel embedding/LM-head/CE（loss diff `3.62e-3`），仍要求 CP/PP/sequence-parallel 关闭；PP/CP MTP 与完整模型长跑仍未验证；fixed-LoRA EP MTP 继续 fail-closed |
+| fixed LoRA | 已实现 | attention/GDN/MLP/shared/routed expert 目标模块 |
+| dynamic multi-LoRA | 已实现子集 | selected v2 默认把不同 rank/target 的租户按 projection-local 最大 rank 补零，保留各租户 `alpha / logical_rank`，在一个 activation batch 内完成一次 forward/backward，再以一次 FinalizeOnly 调用独立更新各租户参数、m/v 和 optimizer clock；入口 preflight 默认合并 6 个布尔状态 collective，未命中的 target 使用零张量且不更新。`QWEN36_HETERO_PADDED_BATCH=0` 保留按 signature 分组和跨组回滚。checkpoint 可独立恢复 heterogeneous rank/alpha/targets；HTTP 可在显式 `allow_aggregate_loss=true` 或 `X-Rustrain-Multi-LoRA-Capability: v1` capability 协商后有界合并兼容且 adapter ID 不相交的并发请求，响应返回 capability version、aggregate scalar、adapter loss 和真实 optimizer step；`expected_steps` 在原生 TP/EP/DP collective 前做全秩乐观并发校验，过期重试 fail closed；dynamic+MTP 已有单轴 DP、TP2 和 TP2xEP2 uneven-source oracle，默认仍保持单请求 dispatch |
+| microbatch accumulation | 已实现子集 | non-final microbatch 只 backward，final microbatch 才 optimizer；FP32 accumulator 存储/聚合，autograd leaf backward 仍为 BF16 |
+| replicated data parallel | 已实现 | logical-step 边界同步 replicated LoRA；EP expert 参数不走该 reduction |
+| expert parallel | 已实现子集 | 默认 routed-output all-reduce；gated variable-split A2A 已验证 fixed-LoRA 和 native dynamic-LoRA data sharding；`QWEN36_EP_A2A_COMPACT_COUNTS=1` 在大 EP world 只将本 rank send row/receive column/validity flags 拷回 host，避免完整 `world×world` metadata D2H，TP2×EP2×DP2 world8 smoke 已通过；GPU-only variable-count dispatch、异步 overlap 和 DeepEP backend 未实现 |
+| tensor parallel | attention + GDN + dense/expert MLP + vocabulary 子集 | full attention 与 GDN 使用 head-aligned ColumnParallel input projection 和 RowParallel output projection；GDN 的 flat QKV/conv 按 `[Q_local|K_local|V_local]` 重排。embedding/LM-head/CE 使用 vocabulary shard；routed/shared expert 使用 gate/up output shard、down input shard，并已在 TP2 x EP2 验证。CLI/server 均可通过共享 v5 rank shard 合并标准 PEFT |
+| sequence parallel | 已实现受限子集 | `QWEN36_SEQUENCE_PARALLEL=1` 仅允许 TP2、dense text-only、projection-aware fixed/dynamic LoRA、CP/EP/DP/PP=1、MTP/aux loss=0；embedding scatter、column all-gather、row reduce-scatter、loss gather 及三租户 selected training 在 H20 两进程 NCCL smoke 中通过，latent-rank dynamic target fail-closed |
+| pipeline parallel | 已实现受限 non-interleaved 子集 | ABI27 stage-local ownership 加上任意 `PP_SIZE>=2` 的固定 shape、单 chunk 1F1B 窗口。新增 opt-in dynamic-LoRA flag：共享 batch 行按注册 tenant 展开，反向恢复 per-tenant token numerator，并复用动态归一化/clip/事务 Adam finalizer；selected tenant 的 ordered request identity 现在在 TP/EP/DP/CP/PP 轴统一校验，并提供 per-tenant loss-report ABI；server 可通过 `QWEN36_PP_MICROBATCHES` 将 `[tenant, seq]` 行批次拆成真实多微批次 1F1B。H20 ABI1 PP2 fused-MLP dynamic smoke 已通过，包含 runtime mismatch、异构 rank 两 tenant/两 microbatch 和独立 loss/Adam 状态。dynamic window 在窗口期间固定 heterogeneous padding，后续本地 shape/mask/tick 错误用 selected-row contract-shape dummy 激活完成固定 P2P，并在 finish 做全 rank fail-closed consensus、禁止 optimizer commit。H20 PP2/PP3 `pp-train-smoke` 均通过固定 LoRA 的正常 loss/gradient/Adam parity、late shape+phase failure 和 divergent-target negative case。仍无 MTP/aux、CP、interleaved/chunked schedule 和 PP-aware reshard |
+| context parallel | dense full-attention ring CP2/CP4 + GDN CP2 | full attention 使用 `QWEN36_CP_FULL_ATTENTION_RING=1` 的 P2P ring、online softmax 和自定义 backward；H20 CP2/CP4 fixed-LoRA smoke 对 CP1 oracle 通过。GDN/hybrid 仍仅 `CP_SIZE=2, TP=EP=DP=PP=1`，拒绝 MTP、aux loss、sequence chunk 和非 grouped dynamic sync；ring 当前串行 comm/compute、未做 zigzag/double-buffer/fused tile，未宣称 Megatron 级吞吐 |
+| distributed checkpoint | 已实现子集 | v5 记录 rank order、完整五维坐标、TP/EP 多轴 placement、fused gate/up segments 及 fixed/dynamic slot identity；任意 multi-rank 拓扑使用 rank 目录。compact receipt 将每 rank preflight 限制为 `O(world_size)` 小 metadata + `O(local state)`；CLI/server 共享 same-topology restore 和标准 PEFT merge。server save/load 使用全 rank prepare/validate/commit-or-abort，load 通过 no-sync shadow context 保证失败不改 live state；fresh destination 以 `RENAME_NOREPLACE` 原子可见。v3/v4 及旧 digest-v5 保持只读兼容；可覆盖 generation/LATEST、断电级 durability、跨 topology reshard 和 PP/CP 未实现 |
+| server tensor transport | 已实现 IPC 子集 | ABI22 parent/worker 共享 64-byte aligned binary slab 与 compact JSON descriptor；worker 不再反序列化三份 `Vec<i64>` 或在 Rust 热路径构造 `tch::Tensor`。普通 tensor HTTP routes 在 decode 前只允许一个 in-flight 请求，忙时快速返回 `429`；`/train_multi` 允许有界并发 admission，并仅对显式 `allow_aggregate_loss=true` 或 `X-Rustrain-Multi-LoRA-Capability: v1` 协商的兼容请求执行短窗口 GPU coalescing，响应声明 capability version 和 per-adapter result contract。带 `expected_steps` 的 optimistic 请求现在保持 request-local failure domain，不参与跨请求 coalescing；无版本 guard 的请求仍走 bounded batching。新增 `/train_multi_binary` 固定版本 little-endian wire，直接解析 int64 tensor sections 后写入 IPC slab，绕过 base64；普通控制面和 checkpoint transaction 使用 bounded FIFO 单消费者 scheduler（默认队列 32，可由 `RUSTRAIN_EP_DISPATCH_QUEUE_CAPACITY` 调整），每个 accepted job 在 `spawn_blocking` 上串行执行。native host-i64 已有 opt-in pinned staging，但仍缺稳定异步 H2D 和跨请求 GPU step overlap |
+
+## 与 Megatron-LM 的关键差距
+
+### 并行语义
+
+Megatron 的通用 MLP 使用 ColumnParallel fc1、local gated activation 和 RowParallel fc2，并进一步提供 fused fc1/activation、sequence parallel、通信重叠和 sharded-state 支持。当前 Qwen native 已补上算法等价的 separate gate/up row shard 与 down column shard，但仍是两次独立 GEMM。
+
+本地 Megatron-LM 的 `experimental/lite/megatron/lite/model/qwen3_5` 已包含直接的 Qwen3.5 实现和 LoRA adapter，而不是只能依赖外部 bridge：`primitive/modules/gqa.py` 使用 fused QKV ColumnParallel 和 O RowParallel；`gated_delta_net.py` 对输入/输出投影和 local heads 做 TP，并接入 sequence parallel、context parallel 与 FLA 高性能路径；`lora.py` 的 `LinearLoRA` 还能拆分 latent rank/output 并配套 gather、reduce-scatter 或 input-gradient all-reduce。没有发现显式的 Qwen3.6 model registration。当前 Qwen native 已补上真实但严格受限的 TP2 sequence-parallel dense、dense GDN CP2 和 dense full-attention CP2/CP4 ring；数学布局与 Megatron 的 gather/reduce-scatter 或 sequence/head exchange 方向一致，但 ring 尚未做 zigzag 负载均衡、双缓冲通信计算重叠或 fused tile softmax。full attention 的 Q/K/V 仍是三次独立 GEMM，LoRA 的 replicated 一侧也比 Megatron Lite 更保守，MoE、PP/CP、多轴组合和生产级 fused backend 仍是主要差距。
+
+PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 attention state 上做跨 rank 通信。ABI26 已把 PP/CP communicator 纳入 `TrainingContext` 和进程级缓存，并通过五轴 size consensus 防止 rank 执行不同的 split 序列；ABI27 补上 stage ownership，当前实现进一步把固定 LoRA 窗口推广到任意 `PP_SIZE>=2` 的固定 shape、单 chunk non-interleaved 1F1B；随后新增了受限 dynamic-LoRA window，PP 版本仍明确拒绝 CP/MTP/aux/interleaving，GDN CP2 dynamic-LoRA 则走独立的 sequence gather/gradient all-reduce 路径。PP control collective 使用独立 communicator，并只在首个微批次建立全局 contract，避免把不同 stage 锁步而与 activation/gradient P2P 互锁；后续本地 shape/mask/tick 错误不直接 reset 单 rank，而是用固定 contract shape 的 dummy 激活完成同样的 P2P 计数，再在 finish 以 control all-reduce fail closed 并清理累积梯度。固定 LoRA 的 H20 PP2/PP3 `pp-train-smoke` 正常 loss、gradient、parameter 和 Adam parity，以及 late shape/phase failure 均通过；H20 ABI1 PP2 dynamic window 也已通过。CP2 x PP2 smoke 仍只验证正交 group 和全网 max 传播；新增的 sequence-parallel 只覆盖 TP2 dense 单轴，不能替代 CP ring attention 或任意五维组合。NCCL 文件 rendezvous 按 generation 隔离：launcher 使用共享 `RUSTRAIN_LAUNCH_OUTPUT_DIR`，直接多节点任务必须提供唯一 `RUSTRAIN_NCCL_RUN_ID` 和 `RUSTRAIN_NCCL_SYNC_DIR`。除严格受限的 dense GDN CP2 外，其余 CP model execution 仍拒绝；MTP/aux 和 interleaved/chunked PP 也仍在读取输入前 fail-closed。
+
+当前 DP/EP 仍不是完整 Megatron 语义：DP 同步 replicated LoRA 梯度并按租户 token count 归一化，expert 参数留在 EP rank；sharded EP 使用 variable-split dispatch/inverse combine 和 fixed/dynamic LoRA data sharding。ABI19 已把所有 top-k assignment 合并为一次 dispatch，并把 local expert 计算合并为 grouped GEMM，但 count planning 仍可见于 host，且没有 fused permutation、异步 overlap 或 DeepEP backend。server 广播同一个 global batch descriptor，worker 按五维 topology 选择本 source rows：TP peers 保持相同，sharded EP 使用 `DP*EP` sources，replicated EP 只按 DP 分片。ABI22 将 parent/worker 数据从每 worker 全量 JSON vector 反序列化改为共享 binary slab；worker 直接借用 slab span 并调用 C++ host-i64 coarse ABI，Rust worker 热路径不再执行 `Tensor::from_slice`、`reshape` 或 `to_device`。普通控制面/检查点现在由 bounded FIFO 单消费者调度器承接，避免 Tokio handler 被阻塞式 IPC 占用；tensor 请求仍单 in-flight，避免多个 48 MiB body 同时驻留，因而它不是多租户 GPU batching。IPC timeout 会永久 poison channel，首次失败立即终止并回收 worker；pre-publish `InvalidInput` 不 poison channel。health 在 terminal failure 后返回 unavailable，partial launch 同样回收已启动 ranks。HTTP ingress 现在同时支持 binary slab wire 与 JSON fallback；native host-i64 有 opt-in pinned staging，coalescing deadline 也从 admission 起算以覆盖前一个 GPU step，但跨请求 H2D/compute overlap 仍未实现，因此仍不能视为最终的高吞吐服务传输。
+
+H20 ABI1 `cp-attention-smoke` 进一步验证了受限 dense full-attention CP2 bridge：local Q 使用 CP rank 的 RoPE offset，normalized K/V 经可微 sequence all-gather 后以显式 global causal mask 做 SDPA，right-padding 跨 CP boundary 的 CP2 与 CP1 eval/loss 差为 `1.117e-3`，Adam `m/v` 差为 `5.595e-6/1.284e-9`；rank-local `QWEN36_CP_FULL_ATTENTION_KV_GATHER` mismatch 在 collective 前 fail-closed。Rust CLI/server 入口仅允许 CP2、TP=EP=DP=PP=1 且显式开启该 flag；这不是 ring attention，也未证明长序列吞吐。
+
+### 优化器与恢复
+
+固定 LoRA 的 Adam 状态可导出/导入，native context 的 logical step 与 checkpoint step 对齐；恢复时校验 canonical base-model path，CPU safetensors state 会复制到 shadow CUDA/FP32 allocation。dynamic adapter 的请求频率不同，每租户拥有独立 optimizer step、m/v 与 FP32 accumulator。server checkpoint load 只有在 fixed/dynamic LoRA、全部 m/v、各 optimizer clock 和 session metadata 都 hydrate 成功后才交换 context。heterogeneous selected v2 默认把不同 signature 的 adapter 安装到同一个 padded registry，一次 TrainOnly 后再执行一次 FinalizeOnly；低 rank 的真实 leaf 通过可微 `cat` 补零，inactive projection 使用同 geometry 的零张量，因此不会破坏 autograd 或错误更新 target hole。旧的 deterministic signature grouping 与 persistent shadow 跨组回滚仍可通过环境变量启用。若二次恢复（registry、梯度清理、rollback 或 restore）失败，native context 现在标记为 poisoned；后续动态请求先在 TP/EP/DP 做健康共识并把全 worker group 一致隔离，EP parent 收到 terminal result 后立即标记 unavailable 并回收全部 worker。H20 world4 `TP2 x EP2` 用 rank0-only 注入验证四 rank 均 fail-closed 且无 collective hang；正常 Adam 注入失败仍恢复为 healthy。同时仍缺少基于模型内容 fingerprint 的身份校验。
+
+### 性能工程
+
+ABI24 adds a packed LoRA gradient synchronization path: FP32 accumulators are
+grouped by EP/DP/TP reduction mask and each populated bucket/axis uses one NCCL
+all-reduce. `QWEN36_PACKED_LORA_SYNC=0` keeps the per-tensor grouped fallback.
+This reduces collective launch count, but it still runs after full backward and
+the token-count CPU fence; it is not Megatron-style backward bucket overlap or
+reduce-scatter. H20 TP2 fixed-LoRA smoke and TP2 x DP2 oracle passed with both
+settings; the synthetic TP2 benchmark was `77.970 ms` packed versus `78.024 ms`
+fallback p50, within measurement noise.
+
+ABI28 also packs the selected-v2 preflight booleans (health, request validity,
+loss-report mode/capacity, and accumulation state) into one per-axis MIN
+all-reduce. On H20 TP2 x EP2 with the synthetic
+heterogeneous workload (`B=8`, `S=128`, eight active tenants), three interleaved
+runs measured p50 old/packed pairs of `13.394/13.256`, `13.219/13.153`, and
+`13.638/13.199 ms`; the median improved by about `1.5%`. An Nsight Systems
+process-tree trace reduced `ncclDevKernel_AllReduce_Sum_u32_RING_LL` instances
+from `1032` to `752`; u64 registry/hash collectives were unchanged. This is a
+control-plane launch reduction, not communication/compute overlap or a matched
+Megatron result.
+
+The GDN backward path now has a separately validated CUDA optimization. The
+reverse recurrence restores `R_t` while reading `S_t` for the direct output
+gradient, so the old standalone state-undo sweep and three per-token barriers
+are removed. Fusion is enabled by default and can be disabled with
+`QWEN36_GDN_RECURRENT_FUSION=0`. On H20, the matched `B=2,S=512,H=2048,L=3`
+benchmark improved single-rank p50 from `80.272` to `75.685 ms` and TP2 from
+`75.908` to `70.885 ms`; TP2 at `B=8` improved from `157.986` to `148.216 ms`.
+An independent ATen recurrence-backward TP2 smoke preserved fixed/dynamic loss,
+FP32 m/v and parameter deltas (`adam_error=0` on both ranks). This is a local
+GDN-kernel improvement, not evidence of FLA, sequence/context parallelism, or
+matched Megatron end-to-end throughput.
+
+Server control-plane and checkpoint commands now use a bounded FIFO
+single-consumer dispatcher. Accepted jobs run one-at-a-time on
+`spawn_blocking`, preserving IPC collective order while keeping Tokio handlers
+responsive; queue pressure returns `429` and scheduler/worker failure returns
+`503`. The default capacity is 32 (hard maximum 4096). Tensor train/eval routes
+still admit only one 48 MiB body and reject while busy, so this is runtime
+backpressure/fairness rather than cross-tenant GPU batching. The separate
+`/train_multi` route now supports the opt-in bounded coalescer described below.
+
+The EP `/train_multi` route now adds an opt-in cross-request coalescer. A client
+may set `allow_aggregate_loss=true`, or advertise
+`X-Rustrain-Multi-LoRA-Capability: v1`; the latter is a versioned opt-in for
+per-adapter loss/step results and the explicit loss scope. Without either
+signal, the legacy one-request dispatch and request-local loss are preserved.
+Opt-in requests are grouped only when
+session, sequence/source layout, and adapter IDs are compatible and disjoint.
+The scheduler keeps up to `RUSTRAIN_MULTI_LORA_BATCH_MAX_OPEN_WINDOWS`
+compatible layout buckets open (default 8, hard maximum 64), so interleaved
+sequence/source layouts no longer seal otherwise mergeable windows early.
+Admission also enforces `RUSTRAIN_MULTI_LORA_BATCH_MAX_RANK_WORK` over
+`adapter_count * max_lora_rank` so heterogeneous rank padding cannot silently
+turn a small request window into an oversized GEMM. The coalescer rebuilds
+source-major rows into one native heterogeneous batch, is bounded by
+request/adapter/rank-work/payload limits, and seals its window before
+ordinary tensor, registry, or checkpoint operations. Responses expose
+`loss_scope=coalesced_batch` and the number of merged requests. ABI25 requires
+the native report symbols and returns adapter-ordered losses: it sums token-loss
+numerators and supervised-token counts only over DP and source-sharded EP,
+without double-counting TP or replicated EP. The server slices that vector back
+to each request while retaining the explicitly labelled aggregate scalar for
+wire compatibility. Older native libraries are rejected at load time; within
+ABI25, report/legacy call mode is negotiated collectively before report-only
+reductions. IPC slab wire version 2 rejects old parent/worker layouts, and the
+coordinator verifies rank-consistent report values. The coalescer is capped at
+2048 adapters for the fixed 256 KiB result slot; an oversized serialized result
+is replaced by a compact error and still signals the waiting parent. This is a
+scheduling/transport optimization, not PP/CP or DeepEP communication overlap.
+The bounded multi-bucket scheduler passed 53/53 server tests in both the local
+ABI1 environment and on H20, including interleaved-layout retention, deadline,
+capacity, binary ingress, and rank-work admission cases.
+
+An opt-in asynchronous A2A metadata prototype was tested against the
+compact-count path on H20. EP4 `H=4096, I=8192, S=128` regressed p50 from
+`5.6632 ms` to `5.8951 ms`; EP8 repeated runs were effectively tied at p50
+(`5.6856` vs `5.6524 ms`) while async p95 worsened to `7.1466` from `5.9853 ms`.
+The uncommitted prototype was removed; `QWEN36_EP_A2A_ASYNC_METADATA` is not a
+supported runtime flag.
+
+The `/train_multi_binary` endpoint accepts version-1 `RLM1` requests with a
+56-byte header, adapter IDs/optional expected steps, and three contiguous
+little-endian int64 tensor sections. It validates exact `[batch, seq]` geometry,
+alignment, counts, and trailing bytes before dispatching the same coalescer and
+native ABI as JSON requests. This removes base64 expansion and JSON tensor
+materialization from the ingress path; the IPC worker still receives the same
+validated slab contract.
+
+Variable-split A2A also has an opt-in compact count transfer. With
+`QWEN36_EP_A2A_COMPACT_COUNTS=1` (enabled by default for communicator worlds
+of at least eight), the host receives only `O(world)` counts instead of the
+complete `world x world` matrix; the variable-count NCCL enqueue remains on the
+host because the current prebuilt NCCL API still requires receiver counts before
+`ncclRecv`. H20 TP2 x EP2 x DP2 world8 `tri-smoke` passed with the default path.
+On the small TP2 x EP2 benchmark the compact path was not enabled by default;
+forced compact A/B was `22.707 ms` p50 versus `22.385 ms` legacy, so EP2/EP4
+keep the legacy path.
+
+The A2A dispatch now packs the int64 token and expert IDs into one two-column
+metadata message per peer, reducing the forward dispatch protocol from three
+NCCL messages (hidden/token/expert) to two (hidden/metadata). The runtime hash
+includes `QWEN36_EP_A2A_PACKED_METADATA`; setting it to `0` restores the split
+protocol. H20 local and TP2 x EP2 distributed smokes passed. On the same
+synthetic dynamic-MoE workload (`B=8,S=128,H=1024,I=2048`, 8 tenants, 50
+iterations), split/packed p50 was `14.155/14.051 ms` and p95 was
+`16.849/16.016 ms`. The roughly `0.7%` p50 change is small enough to treat as
+neutral rather than a stable throughput win. Count planning still crosses the
+host boundary, so this remains materially below DeepEP's GPU-only dispatcher.
+
+GDN fused backward 现在可通过
+`QWEN36_GDN_STATE_CHECKPOINT_STRIDE` 保存固定 token 边界的 FP32 recurrent
+state，并在每个 reverse chunk 开始时从精确 chunk-end state 重启。默认值
+为 `0`，因为每层额外占用约
+`BH * (ceil(S / stride) + 1) * DK * DV * 4` bytes。该路径限制跨 chunk 的
+反向重建误差累积，但 chunk 内仍通过 clamp 后的 decay 反除，不能宣称已
+解决单 token 极小 decay 的数值问题，也没有提供 FLA/sequence-parallel
+chunk kernel 或 packed `cu_seqlens`。
+
+在此 checkpoint 基础上，显式 `QWEN36_GDN_CHUNKWISE_BWD=1` 启用两阶段 AOT
+backward：轻量串行 pass 生成精确 `dS` chunk 边界，随后以
+`(batch-head, chunk)` CTA 并行 replay 完整参数梯度。H20 TP2
+`B=2,S=512,H=2048,L=3,stride=32` 的重复 matched p50 为
+`71.484 -> 53.458 ms`（`-25.2%`），TP2 stride-2 fixed/dynamic oracle 的
+`adam_error=0`。该 flag 和 checkpoint stride 已进入 distributed topology
+hash，非法 checkpoint 配置 fail-fast，checkpoint address 使用 `size_t`。
+该路径不能默认开启：single-rank p50
+`76.195 -> 81.843 ms`（回退 `7.4%`），TP2 `B=8` 仅改善约 `1.15%`，且
+该 workload 的 state checkpoint 约增加 `100 MiB/GPU`。
+
+H20 ABI1 的 TP2 full-attention oracle 已在 `QWEN36_FUSED_QKV` 关闭和开启两种模式下通过：fixed eval/loss 差为 `4.40e-4`，selected dynamic loss 差为 `6.82e-5`，各 Q/K/V/O 参数差不超过 `3.05e-5`。这证明 fused 路径与未融合路径等价，不代表已完成完整模型性能胜出。
+
+full-attention dynamic LoRA 另新增 `QWEN36_FUSED_LORA_QKV_A=1` opt-in：当 Q/K/V 的 adapter layout、rank、batch、input width、device 都兼容时，将三个 A stack 沿 rank 维拼接，只执行一次 `A@x` batched matmul，再切分 latent activation 并执行各自的 B matmul；异构请求自动回退旧路径。latent-rank TP 只复制一次输入，三个 output delta 仍分别执行既有 all-reduce；flag 纳入 distributed runtime hash，rank-local 不一致会 fail-closed。H20 local smoke 和 TP2 fixed/dynamic oracle 均通过。TP2 x EP2 synthetic A/B 在 `B=2,S=128,H=1024,rank=16,2 active tenants` 时 p50 为 legacy/fused `9.354/9.291 ms`（约 `0.7%`）；`B=8,S=512,H=2048,8 active tenants` 的重复结果在约 `20.2-21.4 ms` 间波动，未形成稳定收益。因此默认关闭，当前证据只支持正确、可回滚的实验路径，不支持“高性能已完成”的结论。
+
+MoE shared expert 另新增 `QWEN36_MOE_SHARED_OVERLAP=1` opt-in：每个 `TrainingContext` 懒创建一个稳定的 nonblocking CUDA stream，producer/completion event 将 shared dense expert 与 routed sort/A2A/grouped GEMM 排队到不同 stream，residual combine 前才等待；runtime hash 包含该 flag，context free 会销毁 stream，默认路径不改变。H20 TP2 x EP2 fixed/dynamic smoke 及 rollback/poison 负测通过。交错 H20 synthetic A/B 中，`B=2,S=128,H=1024` 没有稳定收益；`B=8,S=512,H=2048` p50 legacy/overlap 为 `53.195/53.576 ms`，allocator reserved 为 `10.004/12.182 GiB`。由于 overlap 未带来稳定 step-time 改善且增加 workspace/cache 压力，保持默认关闭；这不是 DeepEP/Megatron overlap 等价证明。
+
+本轮新增 `QWEN36_GDN_FUSED_CP_EXCHANGE=1` opt-in：在 CP2 dense GDN 中把 peer-major `[Q|K|V|Z|A|B]` payload 的两次 sequence-to-head exchange 合为一次，autograd backward 也合为一次 inverse exchange；总 payload bytes 不变，配置进入 runtime hash，rank-local flag 不一致会在 P2P 前 fail-closed。H20 `B=2,S=512,H=2048,L=3,rank=8,warmup=5,iters=30` 的 rank0 p50/p90 为 legacy/fused `103.811/105.456 ms` 与 `104.011/107.565 ms`；`B=8` 为 `227.745/231.473 ms` 与 `227.952/236.727 ms`。B=2 fused allocator peak `0.904172/0.996094 GiB`（allocated/reserved），legacy `0.904271/0.990234 GiB`；B=8 fused `2.531722/2.777344 GiB`，legacy `2.531722/2.783203 GiB`。最小 Nsight trace（1 layer、warmup 1、iters 2）确认 `ncclDevKernel_SendRecv` 从 `36` 降至 `24`（每步每层 `6 -> 4`），但端到端 p50 没有稳定改善，故默认关闭；这是 native synthetic CP2 A/B，不是 Megatron/FLA 对照，也没有证明通信计算 overlap。
+
+当前粗粒度 C++ FFI、packed/grouped MoE、GDN head TP、vocabulary TP 和 activation checkpoint/offload 是有效优化；full attention 的 activation-level LoRA 路径新增 `QWEN36_FUSED_QKV=1` opt-in，可将 frozen Q/K/V base projection 合成一次 GEMM，但由于会复制冻结权重，默认关闭，必须在目标 GPU 做显存/延迟 A/B。TP collective 仍同步执行。heterogeneous selected v2 的 padded 路径消除了按 signature 重复完整 trainer 和组间 GPU success fence；在 H20 TP2 x EP2 的四租户 synthetic workload 上，p50 从 grouped 的 `17.750 ms` 降到 `13.841 ms`，unique-token throughput 从 `57.69k/s` 提升到 `73.98k/s`，但 allocator peak 从约 `0.240 GiB` 增至 `0.354 GiB`，且 10 次样本的 p95 为 `20.152 ms`，未优于 grouped 的 `19.441 ms`。这证明 unified padded batch 对当前 native 路径有效；HTTP 仅在显式 `allow_aggregate_loss=true` 时合并兼容请求，且 aggregate loss 会被明确标记，不是 Megatron 对比。尚无 Megatron/Transformer Engine 级别的完整模型端到端数据：没有同一 GPU、序列长度、microbatch、精度和通信配置下的 tokens/s、显存、扩展效率对照，也没有 FP8/FP4 参数与 fused attention/DeepEP 的 Qwen 路径。
+
+新增 `QWEN36_FUSED_MLP_FC1=1` opt-in，把冻结 dense/shared gate/up projection 合成一次 FC1 GEMM，再保留每个租户 LoRA delta 的独立累加；该配置也纳入 PP runtime identity。H20 TP2 x EP2 synthetic dynamic-MoE benchmark（B=2、8 tenants、2 active）中，p50 `11.133 -> 10.672 ms`（约 `4.1%`），但 p95 `12.231 -> 13.777 ms`，allocator peak 约增加 `3.6 MiB/GPU`；B=8、8 active tenants 中 p50 `14.034 -> 13.930 ms`（约 `0.7%`），p95 `15.488 -> 15.706 ms`，allocator peak 约增加 `2.4 MiB/GPU`。因此该路径默认关闭，只作为目标硬件上的可回滚 A/B 选项；native smoke 和 PP2 dynamic isolation 均已通过。
+
+ABI15 synthetic GDN TP benchmark 使用 3 层、S=512、H=2048、16 K heads、32 V heads 和 LoRA rank 8。在 B=2 时 single/TP2 p50 为 `81.10/约 76.06 ms`，只有约 `1.07x`；在 B=8 时为 `303.25/约 158.04 ms`，达到约 `1.92x`，每卡 observed resident 从 `4.59 GiB` 降到 `3.54 GiB`。`nsys` 的 B=2 TP2 trace 中 fused delta-rule backward 占 GPU kernel time `80.3%`、forward 占 `7.4%`，NCCL all-reduce 合计低于 `0.4%`。小 batch 下 single 和 TP2 都只需相近的 persistent-block wave 数，因此不能期待 head TP 自动加速；dynamic multi-LoRA batching 提高 BH 后才接近线性 scaling。这是 synthetic native 结果，不是完整模型或 matched Megatron 对比。
+
+本次 native benchmark 没有证明 gated A2A 的端到端 step-time 优势：在该小型 workload 上 sharded A2A 的中位 step 反而比 legacy 高约 `23%`。它没有实现 DeepEP 的 fused permutation、GPU-only split planning 或通信计算 overlap，也没有覆盖 H=`2048`/E=`256` 的完整 Qwen3.6 workload。legacy 模式复制输入 batch，因此必须同时报告唯一样本吞吐，不能只看所有 rank 的 processed tokens/s。
+
+上述结果是 ABI0 的纯 EP2 小模型历史基线。ABI19 在 TP2 x EP2、相同 source-sharded 语义下做了 matched routing-slot/packed 对照，并在每步训练后沿 EP、expert-DP、TP 依次做 MAX，从而让 4 个 rank 使用相同的全局最慢 rank 样本：p50 `9.007 -> 6.499 ms`，unique-token throughput `56.84k -> 78.78k/s`。Nsight Systems process-tree trace 同时记录到 CUDA kernel launch `6812 -> 5168`、NCCL SendRecv `96 -> 48` 和 AllGather `24 -> 12`；该 trace 采集于 benchmark-only world-max metrics collective 加入之前，因此不包含这个计时区间外的聚合。它证明 packed dispatcher 对当前 native 路径有效，但仍不等于 DeepEP 或 Megatron 的重叠能力。新增 TP2 dense sequence-parallel H20 smoke 只验证 all-gather/reduce-scatter/loss-gather 与固定 LoRA 更新，不提供吞吐结论。
+
+目标 H20 的 ABI1 环境是 CPython 3.12、PyTorch 2.12.1+cu130（ABI=1）、cuDNN 9.20、NCCL 2.29.7、Triton 3.7.1 和 NVSHMEM 3.4.5；机器为 8 张 H20 SM90。Megatron Core、Transformer Engine、FLA、DeepEP、flash-attn 和 Apex 均未安装。PyTorch 自带 SDPA flash backend、`_scaled_mm`、cuBLAS/cuDNN/NCCL，因此仍是当前直接满足约束的生产计算栈。
+
+依赖审计的结论需要区分“预构建产物能运行”和“能无 Python 接入 rustrain”：PrimeIntellect `prime-rl v0.5.0` 提供了可运行的 DeepEP wheel `deep_ep-1.2.1+73b6ea4.cu13-cp312-cp312-linux_x86_64.whl`，SHA256 为 `80369bcbf664d8931950f529e71b549a0e0808c6953b5c8c3dccc91a75770f36`。在 H20 Torch2.12.1/NVSHMEM3.4.5 上，8-GPU `get_dispatch_layout -> dispatch -> combine` 误差为 `0.0`，SM90 检测为真；top1 `4096x7168` BF16 roundtrip p50 约 `0.779 ms`。但 wheel 只有 `PyInit_deep_ep_cpp` 动态业务入口，DeepEP 初始化/IPC 使用 hidden C++ symbols 和 `pybind11::bytearray`，没有稳定 C ABI 或头文件；直接集成需要嵌 CPython/pybind，或者供应方提供预构建 C-ABI shim。自己重编 DeepEP 或 shim 会违背当前“依赖必须预构建、禁止 Python JIT/自构建依赖”的约束，因此本轮只记录为候选验证通过，不接入生产路径。
+
+FlashAttention 另有社区预构建 `flash_attn-2.8.3+cu130torch2.12-cp312` wheel，在当前 H20 可 import 并完成 BF16 GQA forward/backward，forward 约 `0.280 ms`，与 PyTorch SDPA 约 `0.286 ms` 基本相同；由于供应链和 pybind ABI 风险，不能证明值得替换当前 SDPA。TE 2.16.1 的 cp312/cu13/ABI1 wheel 针对 NVIDIA Torch 26.05 的另一 commit；FLA 的 GDN 使用 `triton.jit`/autotune；FA4 依赖 CuTe compile；这些仍不满足当前生产约束。Megatron Core、TE、FLA、DeepEP、flash-attn 和 Apex 仍未安装在目标 Pod。
+
+本地 `/root/code/Megatron-LM` 为 0.19.0（commit `ec2aff43`）。其声明的 GDN/MoE 高性能路径依赖 TE、FLA、FlashAttention/DeepEP，其中锁定环境还包含 git/source build 或 Triton JIT，不能拿到当前 H20 环境直接做合规的 matched run。当前 H20 原语 microbenchmark 只能说明硬件基础可用：SDPA GQA `B=2,Hq=32,Hkv=8,S=1024,D=128` forward 约 `0.286 ms`、forward+backward 约 `1.099 ms`；BF16/FP8 GEMM 约 `132.7/270.6 TFLOP/s`；NCCL A2A 约 `206 GB/s/rank` median。它们不是完整模型 benchmark。因此当前不能诚实地产出 matched Megatron-LoRA tokens/s、显存和扩展效率对照，也没有通过 Python JIT 或自构建依赖绕过限制。
+
+## 验证边界
+
+已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；legacy EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。Replicated A2A 与 fixed-LoRA sharded A2A 也均通过两 rank full-expert reference；sharded token counts `[1,3]` 的加权 loss 与 global reference 相差约 `9.6e-7`，m/v 最大差 `1.22e-5` / `3.92e-9`，Adam oracle 差为 `0`。H20 ABI1 dynamic sharded full-reference smoke 在两 rank 返回 `0`：dynamic grouped-expert 参数最大差 `1.53e-5`，m/v 最大差 `4.88e-5` / `5.75e-8`。ABI13 dense base-MLP TP2 smoke 对 gate/up/down 使用半尺寸本地权重，eval/train loss 与完整权重参考相差 `1.84e-5`。
+
+ABI14 的 4Q/2KV-head GQA full-attention TP2 smoke 覆盖 Q/K/V/O fixed 和 selected dynamic LoRA。fixed eval/loss 与完整参考最大差 `4.40e-4`；Q/K/V-B 与 O-A 梯度最大差分别为 `1.83e-4`、`1.14e-5`、`3.05e-4` 和 `1.83e-4`，FP32 Adam m/v 最大差 `1.53e-5` / `6.63e-9`，本地标准 Adam 公式误差小于 `3.73e-9`。selected dynamic loss 最大差 `6.82e-5`，Q/K/V/O 参数最大差 `3.05e-5`。
+
+H20 ABI1 PP2 native smoke 进一步覆盖两 tenant、两 microbatch、每 tenant 不同 token 权重的 dynamic 1F1B。两 rank 对乱序 selected-adapter 请求一致拒绝；正确请求的 report loss 在两 rank 一致，aggregate loss 与按 tenant token 数加权的 report 相差小于 `1e-5`。选中 tenant 的 q_proj-B 更新而未选 tenant 的参数、Adam m/v 与 step 保持逐位不变；PP2 fixed-LoRA parity 仍通过。追加的异构 rank (2/3) selected 请求通过 `qwen36_add_lora_v2` 注册并完成同一隔离检查。第二个 microbatch 注入 malformed target 后，PP 继续消费 contract-shaped dummy work，fallback batch 保持 selected tenant 数量而不是 registry 总数，最终两 rank 在 finish 一致失败且无 P2P count mismatch。该测试仍限于单 chunk、固定 shape 的 1F1B，不代表 interleaved PP 或 CP model execution。
+
+ABI15 的两层 GDN base-TP smoke 覆盖复合 QKV/conv head shard、Z/A/B/A_log/dt_bias、replicated norm、out-proj input columns，以及五种 GDN projection 的 fixed/selected dynamic LoRA。fixed loss 最大差 `1.30e-4`，rank-local factor 梯度最大差 `2.44e-4`，FP32 Adam m/v 最大差 `3.43e-6` / `2.49e-10`；dynamic loss 最大差 `1.16e-3`，m/v 最大差 `6.10e-6` / `1.22e-9`，标准 Adam 公式误差均为 `0`。latent-rank-only GDN TP2 回归也通过，loss 差 `5.46e-5`。BF16 参数直接对照最多跨约两个量化 bin，因此验收以 FP32 梯度、m/v 和本地 Adam 公式为主。ABI21 world8 TP2 x EP2 x expert-DP2 shadow resume 在 no-sync attach 前故意扰动 DP 非零 rank 的 fixed LoRA，attach 后保持 bitwise 不变；hydrate 后 fixed/dynamic LoRA、Adam m/v 和 tenant clock 精确一致，第二步的参数/m/v 与连续训练零差，step 为 `2`。同一 world8 fixture 已分别覆盖 source-sharded 与 replicated-source A2A；replicated-source 的 fixed 参数/m 最大差 `1.96e-3` / `2.58e-6`，dynamic 参数/m 最大差 `1.99e-3` / `4.16e-5`，两者 Adam 公式误差均为 `0`。ABI22 的单卡 native smoke 验证 borrowed host-i64 eval 与 tensor ABI loss 零差、fixed train Adam oracle 零差，以及 selected dynamic multi-LoRA 仅更新已选 adapter；同一 ABI22 world8 TP2 x EP2 x expert-DP2 `tri-smoke` 也通过 checkpoint resume、selected isolation 与 Adam oracle。heterogeneous restore smoke 在已有 rank-8 expert adapters 后以 no-sync hydration 注册 rank-9 `q_proj` adapter，验证 checkpoint registry 可保留独立 signature。ABI23 进一步覆盖 live heterogeneous registration、tensor/borrowed host-i64 selected v2 dispatch，以及第二组注入失败后参数、Adam m/v 和 tenant clock 的精确回滚。TP2 x EP2 x expert-DP2 world8 使用 rank-3 `q_proj` 与 rank-4 multi-module tenants 验证所有 rank 的成功更新和失败回滚，未出现 collective 顺序分叉。新的 padded 路径在单卡和 TP2 x EP2 smoke 中覆盖 rank-8/rank-9、dense/expert target holes、独立 clock、零 token tenant 与 rollback；显式 grouped fallback 的 TP2 x EP2 smoke 也通过。没有完成完整大模型长时间训练、真实跨节点通信、CP model execution 或与 Megatron-LM 的同条件 benchmark；PP 目前只有固定 LoRA、固定 shape、单 chunk 的 PP2/PP3 smoke。因此“正确”应理解为已覆盖且有直接 oracle 的子集，而不是所有并行配置。
+
+本轮新增 H20 ABI1 CP2 GDN full-reference smoke：`QWEN36_GDN_FUSED_CP_EXCHANGE=0/1` 两种模式均通过 fixed eval/train 与 dynamic selected training。fixed eval/loss 差为 `5.1117e-4`，五种 GDN LoRA projection 的 A/B 参数差为 `1.9531e-3/1.9989e-3`，Adam m/v 差为 `3.0518e-6/1.4198e-10`；dynamic aggregate/per-tenant loss 差均为 `0`，LoRA 参数差 `9.8801e-4`，Adam m/v 差 `1.5259e-6/1.5669e-10`。两个 selected tenant 正常更新，第三个未选 tenant 的参数、m/v 与 optimizer step 均逐位不变；rank-local fused flag mismatch 也在 P2P 前一致 fail-closed。实现将 CP local hidden 先 gather 后计算 CE，再由 autograd/手动 checkpoint backward 切回 local gradient，并在 dynamic grouped gradient sync 中加入 CP all-reduce；非 grouped dynamic CP 继续 fail-closed。
+
+本轮新增 H20 ABI1 `mtp-dynamic-smoke` 和 `mtp-dp-smoke`：注册三个 dense LoRA tenant，选择两个不同主/MTP token-count 的租户，安装一层真实 frozen MTP prediction layer；单卡 dynamic per-tenant report 与单租户 oracle 的 loss 和参数差均为 `0`，DP2 complementary-source rows 对两个 singleton oracle 的 loss、参数、Adam m/v 差均为 `0`，第三个未选择租户参数/m/v/step 逐位不变。DP2 MTP objective 的 loss effect 为 `1.0821`、parameter effect 为 `2.0142e-3`，证明 MTP 分支实际参与训练。实现一次性 all-reduce `[main_counts, mtp_counts]`（DP>1 时）并将 MTP hidden gradient/report numerator 转到主 loss denominator；fixed micro-step 使用同一比例，seq<3 和 `TP/CP/EP/PP>1` 组合 fail-closed。
+
+本轮再新增 H20 ABI1 `mtp-tp-smoke`：TP2 两 rank 对主层和 dense MTP prediction layer 同时使用 attention/MLP 列/行并行，embedding/LM-head 保持 replicated，两个 selected dynamic tenant 与一个未选 tenant 对照单卡 full-weight oracle。两 rank 的 `loss_diff` 最大 `3.53e-3`、参数差最大 `1.87e-9`、Adam m/v 差最大 `2.14e-5/3.23e-9`，未选状态逐位不变；开启 `QWEN36_FUSED_QKV=1` 与 `QWEN36_FUSED_MLP_FC1=1` 后仍通过。后续真实 MoE prediction-layer TP2 MTP oracle 的 `loss_diff=1.50e-3`、参数差 `3.05e-5`；vocabulary-parallel MTP 的 `loss_diff=3.62e-3`、参数差 `2.01e-3`，未选状态均逐位不变。MTP token counts 在 TP 只做 min/max 一致性检查，不做求和，避免重复放大 hidden gradient。TP2xEP2 uneven source-token-count MTP 也已通过（main/MTP tokens `7/5`，loss diff `4.38e-4`，参数差 `2.01e-3`）。这些仍是受限 oracle，不覆盖 PP/CP MTP 或完整模型长跑。
+
+本轮对 heterogeneous selected-v2 的 registry 管理做了小范围优化：恢复阶段按保存的 canonical index 原位放回 selected adapter，并维护持久排序的 `adapter_id -> canonical slot` 索引，供 distributed validator、selected-v2、expected-step 和 selected-eval 使用；临时 selected/chunk registry detach 时显式使索引失效，canonical vector 恢复后再启用。H20 TP2 x EP2 synthetic dynamic-MoE（`B=8,S=128,H=1024,I=2048`、rank 16、q/expert targets、50 iterations）中，8 registered/8 active tenants 的 p50/p95 为 `13.859/17.561 ms`，1024 registered/8 active 为 `14.390/19.579 ms`，p50 增加约 `3.8%`；allocator reserved 从约 `1.04` 增至 `9.46 GiB/GPU`。local smoke 和 TP2 x EP2 distributed smoke 均通过。该结果只证明 registry scaling 的查找开销受控，不等于完整模型端到端收益；大注册表的 adapter 状态显存仍是容量约束。
+
+本轮进一步将动态 Adam 状态改为默认懒 materialize（`QWEN36_LAZY_DYNAMIC_OPTIMIZER_STATE=1`）。冷租户只保留 live LoRA 与 FP32 gradient slab，首个有全局 token 的 finalizer 才分配 `m/v` 与 transactional shadow；getter/checkpoint hydration 只分配 `m/v`，避免恢复历史冷租户时复制 shadow。step 0 checkpoint 使用 `optimizer_count=0` 表示隐式零状态，同时兼容旧的 full-state step 0 manifest。H20 TP2 x EP2 matched synthetic A/B（`B=8,S=128,H=1024,I=2048`、rank 16、8 active）中，8 registered/8 active reserved `1.035 GiB/GPU`，1024 registered/8 active `2.918-3.148 GiB/GPU`，此前完整状态约 `9.46 GiB/GPU`；p50 `14.017/14.058 ms`，处于噪声范围。native local/distributed smoke 和 ABI1 `cargo check --lib` 均通过，materialization allocation failure 通过现有 collective gate fail-closed。
+
+GDN persistent forward 与 checkpoint replay 的 token-loop barrier 已按列独占数据布局移除，初始化及 reverse reduction barrier 保留。H20 TP2 smoke、chunkwise smoke 和 backward parity 均通过；matched `B=2,S=512,H=2048,L=3` 的 stride32 replay p50 baseline/optimized 为 `106.26/106.73 ms`，无 checkpoint 默认路径为 `102.62/102.34 ms`。该结果只支持低风险的同步微优化，不支持宣称有稳定的多百分比端到端收益。
+
+这里“尚未完成 CP model execution”指任意五维组合、ring attention 和长序列 CP；当前新增的 full-attention CP2 仅是 opt-in KV all-gather correctness bridge，和 GDN CP2 一样不应外推为完整 Megatron CP。
+
+本轮 EP dispatch 新增默认关闭的 `QWEN36_EP_A2A_COUNT_OVERLAP=1`：count
+all-gather、compact metadata copy 和 host prefix bookkeeping 使用独立的
+non-blocking CUDA stream，并与当前 stream 的 token packing 建立 event 依赖；
+flag 关闭时保留原路径。H20 TP2xEP2 dynamic matched smoke（packed A2A、
+fused local permutation/weighted unpermute）以及 local/distributed correctness
+均通过。`B=8,S=128,H=1024,I=2048` 的 50-iteration A/B mean 为
+`13.625/13.807 ms`（overlap on/off），但 `B=8,S=512,H=2048,I=4096` 的
+matched mean 为 `52.758/52.772 ms`，因此没有稳定端到端收益，不能宣称已达到
+Megatron/DeepEP 的通信计算重叠。
+
+H20 依赖审计确认项目 ABI1 PyTorch 2.12.1/cu130 环境没有可直接复用的
+Transformer Engine、DeepEP、FLA 或 standalone flash-attn 预编译包；PyTorch
+自带 Flash-SDPA 在 H20 上可用。没有安装依赖、没有 JIT workaround，当前 native
+ATen/C++ kernel 路径因此仍是可复现基线。
+
+动态 Adam checkpoint 保存新增 host-authoritative CPU snapshot（ABI31）。
+第一次导出按 tenant optimizer clock 将完整 FP32 m/v 从 GPU 拷到临时 host map，
+原子发布后四个 getter 复用该 snapshot；step 0 仍使用隐式零状态，restore 路径
+现在通过 bulk host import 保持 restore 后 GPU resident count 为 0，首次选中 tenant
+时才按 pair hydrate。H20 local 与 TP2xEP2 distributed smoke 验证了 step-0/step-1、
+重复 snapshot、partial-import 原子失败、恢复 parity 和 resident count；这减少保存
+与恢复时的 GPU materialization。新增的动态 Adam host paging 保持默认关闭；仅当
+`QWEN36_DYNAMIC_ADAM_HOST_PAGING=1` 且
+`QWEN36_DYNAMIC_ADAM_RESIDENT_BYTES=<bytes>` 有效时，selected finalizer 全秩成功、
+canonical registry 恢复且 transaction/shadow 释放后，才按 tenant 最近成功访问顺序
+逐出本轮未选中的冷 m/v。逐出先原子发布完整 host snapshot，再用 map swap 清空
+device state；刚训练的 tenant 始终受保护，因此 budget 是 best-effort。默认路径不做
+device 同步，超预算逐出才付出 checkpoint 级同步与 D2H copy；paging 失败保留原 device
+state，也不会把已 commit 的 step 误报为训练失败。native TP/EP smoke 覆盖两 tenant、
+零预算逐出、CPU getter 不触发 hydrate，以及被逐出 tenant 下一步相对 resident oracle
+的 loss、参数和 m/v 精确 parity。该实现尚未覆盖异步 pinned D2H、prefetch、PP window
+逐出或跨进程共享 host tier，不能等同于 Megatron distributed optimizer/offload。
+
+## 继续达到 Megatron 级别所需的最小工作包
+
+1. 为已实现的 opt-in server batch coalescing 增加 HTTP response capability/version negotiation，使支持 per-adapter loss 的客户端可取消显式 aggregate opt-in；补齐 padding rank inflation 容量模型和长尾 benchmark。
+2. 为现有任意 PP size 的受限窗口增加 interleaved/chunked scheduler、dynamic LoRA、PP-aware checkpoint/loader，并为 CP 实现 ring attention/state exchange。
+3. 将 EP dispatch/combine 替换为 fused/异步路径，并测量通信与计算重叠。
+4. 为 dense/expert MLP 增加 fused gate/up FC1 和 MTP 支持；为 GDN 消除 chunk 内 decay 反除并增加 sequence/context parallel 与 packed `cu_seqlens`，为 full attention 融合 QKV/SDPA 并增加 sequence parallel。
+5. 为 server 增加 raw binary HTTP ingress、pinned staging、异步 H2D/request overlap；为 checkpoint 增加 generation/`LATEST`、断电级 durability、跨 topology reshard、可恢复的 pending accumulation state 和旧 v3 attention checkpoint 离线迁移工具。
+6. 在固定硬件和 workload 上，与 Megatron-LM 记录 tokens/s、step time、峰值显存、通信占比和 loss 曲线。
+
+## Native EP Benchmark Artifact
+
+`crates/rustrain-qwen3-6/tests/native_ep_bench.cpp` is a dependency-free
+synthetic baseline for the native Qwen C ABI. It times complete
+`qwen36_train_step` calls (forward, backward, and Adam) with CUDA events and
+accepts `BENCH_SEQ`, `BENCH_HIDDEN`, `BENCH_EXPERTS`, `BENCH_INTERMEDIATE`,
+`BENCH_WARMUP`, and `BENCH_ITERS`. It reports processed and unique tokens/s,
+per-rank step statistics, and free memory. This is not a Megatron-LM
+comparison and does not claim DeepEP or Transformer Engine parity.
+
+Fresh H20 ABI0 run (`seq=128, hidden=256, experts=8, intermediate=256,
+warmup=2, iters=10`, two ranks):
+
+| Mode | Rank | Median step | Mean step | Processed tokens/s | Unique tokens/s | Status |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| Legacy EP (`QWEN36_EP_A2A=0`) | 0 | 5.4752 ms | 5.4902 ms | 46,391 | 23,196 | 0 |
+| Legacy EP (`QWEN36_EP_A2A=0`) | 1 | 5.4672 ms | 5.4591 ms | 46,459 | 23,230 | 0 |
+| Sharded A2A (`QWEN36_EP_A2A=1`, `QWEN36_EP_A2A_SHARDED=1`) | 0 | 6.7192 ms | 6.7459 ms | 37,802 | 37,802 | 0 |
+| Sharded A2A (`QWEN36_EP_A2A=1`, `QWEN36_EP_A2A_SHARDED=1`) | 1 | 6.7261 ms | 6.7458 ms | 37,763 | 37,763 | 0 |
+
+Each rank processed 127 local tokens and 254 processed tokens per step. In
+replicated legacy EP, unique tokens/s is 127 tokens divided by step time;
+sharded A2A has 254 unique global tokens per step. These are synthetic native
+measurements only; workload, precision, packing, and communication semantics
+are not matched to Megatron-LM.
diff --git a/scripts/run_qwen36_native_gdn_tp.sh b/scripts/run_qwen36_native_gdn_tp.sh
new file mode 100755
index 00000000..dd5b8979
--- /dev/null
+++ b/scripts/run_qwen36_native_gdn_tp.sh
@@ -0,0 +1,221 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+usage() {
+    echo "usage: $0 {smoke|smoke-chunkwise|tpdp-smoke|bench-single|bench-tp2|bench-cp2}" >&2
+    exit 2
+}
+
+mode="${1:-}"
+case "$mode" in
+    smoke|smoke-chunkwise|tpdp-smoke|bench-single|bench-tp2|bench-cp2) ;;
+    *) usage ;;
+esac
+
+repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$repo_root"
+
+python_bin="${PYTHON:-python3}"
+if [[ -n "${TORCH_INCLUDE_PATH:-}" && -n "${TORCH_LIB_PATH:-}" && -n "${GLIBCXX_USE_CXX11_ABI:-}" ]]; then
+    torch_include="$TORCH_INCLUDE_PATH"
+    torch_lib="$TORCH_LIB_PATH"
+    cxx11_abi="$GLIBCXX_USE_CXX11_ABI"
+    site_packages="${PYTHON_SITE_PACKAGES:-${torch_include%/torch/include}}"
+    torch_version="${TORCH_VERSION:-prebuilt}"
+else
+    readarray -t torch_config < <(
+        "$python_bin" - <<'PY'
+import pathlib
+import torch
+
+root = pathlib.Path(torch.__file__).resolve().parent
+print(root / "include")
+print(root / "lib")
+print(int(torch._C._GLIBCXX_USE_CXX11_ABI))
+print(root.parent)
+print(torch.__version__)
+PY
+    )
+    torch_include="${torch_config[0]}"
+    torch_lib="${torch_config[1]}"
+    cxx11_abi="${torch_config[2]}"
+    site_packages="${torch_config[3]}"
+    torch_version="${torch_config[4]}"
+fi
+
+cuda_home="${CUDA_HOME:-/usr/local/cuda}"
+cuda_include="${CUDA_INCLUDE_PATH:-$cuda_home/include}"
+nccl_include="${NCCL_INCLUDE_PATH:-$site_packages/nvidia/nccl/include}"
+nccl_lib="${NCCL_LIB_PATH:-$site_packages/nvidia/nccl/lib}"
+
+for required in \
+    "$torch_include/ATen/ATen.h" \
+    "$torch_lib/libtorch.so" \
+    "$cuda_include/cuda_runtime.h" \
+    "$nccl_include/nccl.h"; do
+    if [[ ! -e "$required" ]]; then
+        echo "required prebuilt dependency file not found: $required" >&2
+        exit 1
+    fi
+done
+
+export LIBTORCH_USE_PYTORCH=1
+export LIBTORCH_BYPASS_VERSION_CHECK=1
+export TORCH_INCLUDE_PATH="$torch_include"
+export TORCH_LIB_PATH="$torch_lib"
+export CUDA_INCLUDE_PATH="$cuda_include"
+export NCCL_INCLUDE_PATH="$nccl_include"
+export NCCL_LIB_PATH="$nccl_lib"
+export GLIBCXX_USE_CXX11_ABI="$cxx11_abi"
+
+if [[ -e "$nccl_lib/libnccl.so" ]]; then
+    nccl_link="-lnccl"
+    nccl_file="$nccl_lib/libnccl.so"
+elif [[ -e "$nccl_lib/libnccl.so.2" ]]; then
+    nccl_link="-l:libnccl.so.2"
+    nccl_file="$nccl_lib/libnccl.so.2"
+else
+    echo "prebuilt NCCL library not found under $nccl_lib" >&2
+    exit 1
+fi
+
+nvcc="$cuda_home/bin/nvcc"
+if [[ ! -x "$nvcc" ]]; then
+    echo "CUDA compiler not found: $nvcc" >&2
+    exit 1
+fi
+
+for tool in g++ sha256sum stat; do
+    if ! command -v "$tool" >/dev/null; then
+        echo "required build tool not found: $tool" >&2
+        exit 1
+    fi
+done
+
+fingerprint="$({
+    printf '%s\n' \
+        "python=$python_bin" \
+        "torch_version=$torch_version" \
+        "torch_include=$torch_include" \
+        "torch_lib=$torch_lib" \
+        "cxx11_abi=$cxx11_abi" \
+        "cuda_home=$cuda_home" \
+        "cuda_include=$cuda_include" \
+        "nccl_include=$nccl_include" \
+        "nccl_lib=$nccl_lib" \
+        "nccl_link=$nccl_link"
+    g++ --version
+    "$nvcc" --version
+    sha256sum "$0"
+    stat -Lc '%n:%s:%Y' \
+        "$torch_lib/libtorch.so" \
+        "$torch_lib/libtorch_cuda.so" \
+        "$torch_lib/libc10.so" \
+        "$nccl_file" \
+        "$cuda_home/lib64/libcudart.so"
+} | sha256sum | cut -c1-20)"
+native_root="${NATIVE_BUILD_DIR:-target/native-qwen36-gdn}"
+native_dir="$native_root/$fingerprint"
+mkdir -p "$native_dir"
+
+kernel_lib="$native_dir/libqwen36_kernels.so"
+kernel_sources=(
+    crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp
+    crates/rustrain-qwen3-6/kernels/delta_rule.cu
+    crates/rustrain-qwen3-6/kernels/delta_rule.cuh
+    crates/rustrain-qwen3-6/kernels/fused_kernels.cu
+)
+rebuild_kernel=false
+if [[ ! -e "$kernel_lib" ]]; then
+    rebuild_kernel=true
+else
+    for source in "${kernel_sources[@]}"; do
+        if [[ "$source" -nt "$kernel_lib" ]]; then
+            rebuild_kernel=true
+            break
+        fi
+    done
+fi
+
+if [[ "$rebuild_kernel" == true ]]; then
+    cuda_objects=()
+    for source in delta_rule fused_kernels; do
+        object="$native_dir/$source.o"
+        "$nvcc" -c "crates/rustrain-qwen3-6/kernels/$source.cu" -o "$object" \
+            -O2 -std=c++17 "-D_GLIBCXX_USE_CXX11_ABI=$cxx11_abi" \
+            "-I$torch_include" "-I$cuda_include" -Xcompiler -fPIC
+        cuda_objects+=("$object")
+    done
+
+    g++ -shared -fPIC -std=c++17 -O2 \
+        "-D_GLIBCXX_USE_CXX11_ABI=$cxx11_abi" \
+        crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp \
+        "${cuda_objects[@]}" -o "$kernel_lib" \
+        "-I$torch_include" "-I$cuda_include" "-I$nccl_include" \
+        "-L$torch_lib" "-L$nccl_lib" "-L$cuda_home/lib64" \
+        "-Wl,-rpath,$torch_lib" "-Wl,-rpath,$nccl_lib" \
+        "-Wl,-rpath,$cuda_home/lib64" -Wl,--no-as-needed \
+        -ltorch -ltorch_cuda -ltorch_cpu -lc10 -lc10_cuda -lcudart "$nccl_link"
+fi
+kernel_dir="$native_dir"
+
+common_flags=(
+    -std=c++17 -O2 "-D_GLIBCXX_USE_CXX11_ABI=$cxx11_abi"
+    "-I$torch_include" "-I$cuda_include" "-I$nccl_include"
+    "-L$kernel_dir" "-L$torch_lib" "-L$nccl_lib" "-L$cuda_home/lib64"
+    "-Wl,-rpath,$kernel_dir" "-Wl,-rpath,$torch_lib"
+    "-Wl,-rpath,$nccl_lib" "-Wl,-rpath,$cuda_home/lib64"
+    -Wl,--no-as-needed -Wl,--allow-shlib-undefined
+    -lqwen36_kernels -ltorch -ltorch_cuda -ltorch_cpu -lc10 -lc10_cuda
+    -lcudart "$nccl_link"
+)
+
+smoke_bin="$native_dir/native_tp_gdn_smoke"
+tpdp_smoke_bin="$native_dir/native_tp_dp_smoke"
+bench_bin="$native_dir/native_tp_gdn_bench"
+smoke_source=crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp
+tpdp_smoke_source=crates/rustrain-qwen3-6/tests/native_tp_dp_smoke.cpp
+bench_source=crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp
+if [[ ! -e "$smoke_bin" || "$smoke_source" -nt "$smoke_bin" || "$kernel_lib" -nt "$smoke_bin" ]]; then
+    g++ "$smoke_source" -o "$smoke_bin" "${common_flags[@]}"
+fi
+if [[ ! -e "$tpdp_smoke_bin" || "$tpdp_smoke_source" -nt "$tpdp_smoke_bin" || "$kernel_lib" -nt "$tpdp_smoke_bin" ]]; then
+    g++ "$tpdp_smoke_source" -o "$tpdp_smoke_bin" "${common_flags[@]}"
+fi
+if [[ ! -e "$bench_bin" || "$bench_source" -nt "$bench_bin" || "$kernel_lib" -nt "$bench_bin" ]]; then
+    g++ "$bench_source" -o "$bench_bin" "${common_flags[@]}"
+fi
+
+cu13_lib="$site_packages/nvidia/cu13/lib"
+export LD_LIBRARY_PATH="$kernel_dir:$torch_lib:$nccl_lib:$cuda_home/lib64:$cu13_lib:${LD_LIBRARY_PATH:-}"
+export RUSTRAIN_NCCL_RUN_ID="${RUSTRAIN_NCCL_RUN_ID:-qwen36-gdn-$$}"
+
+case "$mode" in
+    smoke)
+        TP_SIZE=2 "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=2 --no-python "$smoke_bin"
+        ;;
+    smoke-chunkwise)
+        QWEN36_GDN_STATE_CHECKPOINT_STRIDE=2 \
+        QWEN36_GDN_CHUNKWISE_BWD=1 \
+        TP_SIZE=2 "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=2 --no-python "$smoke_bin"
+        ;;
+    tpdp-smoke)
+        TP_SIZE=2 DP_SIZE=2 RUSTRAIN_DATA_PARALLEL=1 \
+            "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=4 --no-python "$tpdp_smoke_bin"
+        ;;
+    bench-single)
+        BENCH_MODE=single WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 "$bench_bin"
+        ;;
+    bench-tp2)
+        BENCH_MODE=tp2 TP_SIZE=2 "$python_bin" -m torch.distributed.run \
+            --standalone --nnodes=1 --nproc-per-node=2 --no-python "$bench_bin"
+        ;;
+    bench-cp2)
+        BENCH_MODE=cp2 TP_SIZE=1 CP_SIZE=2 EP_SIZE=1 DP_SIZE=1 PP_SIZE=1 \
+            "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=2 --no-python "$bench_bin"
+        ;;
+esac
diff --git a/scripts/run_qwen36_native_tp_ep.sh b/scripts/run_qwen36_native_tp_ep.sh
new file mode 100755
index 00000000..325b6f72
--- /dev/null
+++ b/scripts/run_qwen36_native_tp_ep.sh
@@ -0,0 +1,315 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+mode="${1:-smoke}"
+case "$mode" in
+    local-smoke)
+        test_name=native_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_smoke.cpp
+        ;;
+    smoke|gpu-metadata-smoke|tri-smoke|tri-replicated-smoke)
+        test_name=native_tp_ep_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_tp_ep_smoke.cpp
+        ;;
+    sequence-parallel-smoke)
+        test_name=native_sequence_parallel_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_sequence_parallel_smoke.cpp
+        ;;
+    tp-attention-smoke)
+        test_name=native_tp_attention_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_tp_attention_smoke.cpp
+        ;;
+    mtp-dynamic-smoke)
+        test_name=native_mtp_dynamic_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_mtp_dynamic_smoke.cpp
+        ;;
+    mtp-dp-smoke)
+        test_name=native_mtp_dp_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_mtp_dp_smoke.cpp
+        ;;
+    mtp-tp-smoke)
+        test_name=native_mtp_tp_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_mtp_tp_smoke.cpp
+        ;;
+    mtp-tp-ep-smoke)
+        test_name=native_mtp_tp_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_mtp_tp_smoke.cpp
+        ;;
+    cp-gdn-smoke)
+        test_name=native_cp_gdn_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_cp_gdn_smoke.cpp
+        ;;
+    cp-attention-smoke)
+        test_name=native_cp_attention_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_cp_attention_smoke.cpp
+        ;;
+    ep-smoke)
+        test_name=native_ep_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp
+        ;;
+    ep-bench)
+        test_name=native_ep_bench
+        test_source=crates/rustrain-qwen3-6/tests/native_ep_bench.cpp
+        ;;
+    bench)
+        test_name=native_tp_ep_bench
+        test_source=crates/rustrain-qwen3-6/tests/native_tp_ep_bench.cpp
+        ;;
+    pp-cp-comm-smoke)
+        test_name=native_pp_cp_comm_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_pp_cp_comm_smoke.cpp
+        ;;
+    pp-train-smoke)
+        test_name=native_pp_train_smoke
+        test_source=crates/rustrain-qwen3-6/tests/native_pp_train_smoke.cpp
+        ;;
+    *)
+    echo "usage: $0 [local-smoke|smoke|sequence-parallel-smoke|tp-attention-smoke|cp-attention-smoke|mtp-dynamic-smoke|mtp-dp-smoke|mtp-tp-smoke|mtp-tp-ep-smoke|cp-gdn-smoke|gpu-metadata-smoke|tri-smoke|tri-replicated-smoke|ep-smoke|ep-bench|bench|pp-cp-comm-smoke|pp-train-smoke]" >&2
+        exit 2
+        ;;
+esac
+
+repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$repo_root"
+
+python_bin="${PYTHON:-python3}"
+if [[ -n "${TORCH_INCLUDE_PATH:-}" && -n "${TORCH_LIB_PATH:-}" && -n "${GLIBCXX_USE_CXX11_ABI:-}" ]]; then
+    torch_include="$TORCH_INCLUDE_PATH"
+    torch_lib="$TORCH_LIB_PATH"
+    cxx11_abi="$GLIBCXX_USE_CXX11_ABI"
+    site_packages="${PYTHON_SITE_PACKAGES:-${torch_include%/torch/include}}"
+    torch_version="${TORCH_VERSION:-prebuilt}"
+else
+    readarray -t torch_config < <(
+        "$python_bin" - <<'PY'
+import pathlib
+import torch
+
+root = pathlib.Path(torch.__file__).resolve().parent
+print(root / "include")
+print(root / "lib")
+print(int(torch._C._GLIBCXX_USE_CXX11_ABI))
+print(root.parent)
+print(torch.__version__)
+PY
+    )
+    torch_include="${torch_config[0]}"
+    torch_lib="${torch_config[1]}"
+    cxx11_abi="${torch_config[2]}"
+    site_packages="${torch_config[3]}"
+    torch_version="${torch_config[4]}"
+fi
+
+cuda_home="${CUDA_HOME:-/usr/local/cuda}"
+cuda_include="${CUDA_INCLUDE_PATH:-$cuda_home/include}"
+nccl_include="${NCCL_INCLUDE_PATH:-$site_packages/nvidia/nccl/include}"
+nccl_lib="${NCCL_LIB_PATH:-$site_packages/nvidia/nccl/lib}"
+
+for required in \
+    "$torch_include/ATen/ATen.h" \
+    "$torch_lib/libtorch.so" \
+    "$cuda_include/cuda_runtime.h" \
+    "$nccl_include/nccl.h"; do
+    if [[ ! -e "$required" ]]; then
+        echo "required prebuilt dependency file not found: $required" >&2
+        exit 1
+    fi
+done
+
+if [[ -e "$nccl_lib/libnccl.so" ]]; then
+    nccl_link="-lnccl"
+    nccl_file="$nccl_lib/libnccl.so"
+elif [[ -e "$nccl_lib/libnccl.so.2" ]]; then
+    nccl_link="-l:libnccl.so.2"
+    nccl_file="$nccl_lib/libnccl.so.2"
+else
+    echo "prebuilt NCCL library not found under $nccl_lib" >&2
+    exit 1
+fi
+
+nvcc="$cuda_home/bin/nvcc"
+for tool in "$nvcc" g++ sha256sum stat; do
+    if ! command -v "$tool" >/dev/null 2>&1; then
+        echo "required build tool not found: $tool" >&2
+        exit 1
+    fi
+done
+
+export LIBTORCH_USE_PYTORCH=1
+export LIBTORCH_BYPASS_VERSION_CHECK=1
+export TORCH_INCLUDE_PATH="$torch_include"
+export TORCH_LIB_PATH="$torch_lib"
+export CUDA_INCLUDE_PATH="$cuda_include"
+export NCCL_INCLUDE_PATH="$nccl_include"
+export NCCL_LIB_PATH="$nccl_lib"
+export GLIBCXX_USE_CXX11_ABI="$cxx11_abi"
+
+fingerprint="$({
+    printf '%s\n' \
+        "python=$python_bin" \
+        "torch_version=$torch_version" \
+        "torch_include=$torch_include" \
+        "torch_lib=$torch_lib" \
+        "cxx11_abi=$cxx11_abi" \
+        "cuda_home=$cuda_home" \
+        "nccl_include=$nccl_include" \
+        "nccl_lib=$nccl_lib" \
+        "nccl_link=$nccl_link"
+    g++ --version
+    "$nvcc" --version
+    sha256sum "$0"
+    stat -Lc '%n:%s:%Y' \
+        "$torch_lib/libtorch.so" \
+        "$torch_lib/libtorch_cuda.so" \
+        "$torch_lib/libc10.so" \
+        "$nccl_file" \
+        "$cuda_home/lib64/libcudart.so"
+} | sha256sum | cut -c1-20)"
+
+native_root="${NATIVE_BUILD_DIR:-target/native-qwen36-tp-ep}"
+native_dir="$native_root/$fingerprint"
+mkdir -p "$native_dir"
+kernel_lib="$native_dir/libqwen36_kernels.so"
+kernel_sources=(
+    crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp
+    crates/rustrain-qwen3-6/kernels/delta_rule.cu
+    crates/rustrain-qwen3-6/kernels/delta_rule.cuh
+    crates/rustrain-qwen3-6/kernels/fused_kernels.cu
+)
+
+rebuild_kernel=false
+if [[ ! -e "$kernel_lib" ]]; then
+    rebuild_kernel=true
+else
+    for source in "${kernel_sources[@]}"; do
+        if [[ "$source" -nt "$kernel_lib" ]]; then
+            rebuild_kernel=true
+            break
+        fi
+    done
+fi
+
+if [[ "$rebuild_kernel" == true ]]; then
+    cuda_objects=()
+    for source in delta_rule fused_kernels; do
+        object="$native_dir/$source.o"
+        "$nvcc" -c "crates/rustrain-qwen3-6/kernels/$source.cu" -o "$object" \
+            -O2 -std=c++17 "-D_GLIBCXX_USE_CXX11_ABI=$cxx11_abi" \
+            "-I$torch_include" "-I$cuda_include" -Xcompiler -fPIC
+        cuda_objects+=("$object")
+    done
+    g++ -shared -fPIC -std=c++17 -O2 \
+        "-D_GLIBCXX_USE_CXX11_ABI=$cxx11_abi" \
+        crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp \
+        "${cuda_objects[@]}" -o "$kernel_lib" \
+        "-I$torch_include" "-I$cuda_include" "-I$nccl_include" \
+        "-L$torch_lib" "-L$nccl_lib" "-L$cuda_home/lib64" \
+        "-Wl,-rpath,$torch_lib" "-Wl,-rpath,$nccl_lib" \
+        "-Wl,-rpath,$cuda_home/lib64" -Wl,--no-as-needed \
+        -ltorch -ltorch_cuda -ltorch_cpu -lc10 -lc10_cuda -lcudart "$nccl_link"
+fi
+
+test_bin="$native_dir/$test_name"
+if [[ ! -e "$test_bin" || "$test_source" -nt "$test_bin" || "$kernel_lib" -nt "$test_bin" ]]; then
+    g++ "$test_source" -o "$test_bin" \
+        -std=c++17 -O2 "-D_GLIBCXX_USE_CXX11_ABI=$cxx11_abi" \
+        "-I$torch_include" "-I$cuda_include" "-I$nccl_include" \
+        "-L$native_dir" "-L$torch_lib" "-L$nccl_lib" "-L$cuda_home/lib64" \
+        "-Wl,-rpath,$native_dir" "-Wl,-rpath,$torch_lib" \
+        "-Wl,-rpath,$nccl_lib" "-Wl,-rpath,$cuda_home/lib64" \
+        -Wl,--no-as-needed -Wl,--allow-shlib-undefined \
+        -lqwen36_kernels -ltorch -ltorch_cuda -ltorch_cpu -lc10 -lc10_cuda \
+        -lcudart "$nccl_link"
+fi
+
+cu13_lib="$site_packages/nvidia/cu13/lib"
+export LD_LIBRARY_PATH="$native_dir:$torch_lib:$nccl_lib:$cuda_home/lib64:$cu13_lib:${LD_LIBRARY_PATH:-}"
+if [[ -z "${RUSTRAIN_NCCL_RUN_ID:-}" &&
+      ( -z "${RUSTRAIN_RUN_ID:-}" || -z "${RUSTRAIN_ATTEMPT_ID:-}" ) ]]; then
+    export RUSTRAIN_NCCL_RUN_ID="qwen36-tp-ep-$$"
+fi
+
+if [[ "$mode" == "local-smoke" ]]; then
+    WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 TP_SIZE=1 EP_SIZE=1 DP_SIZE=1 \
+        "$test_bin"
+elif [[ "$mode" == "sequence-parallel-smoke" ]]; then
+    QWEN36_SEQUENCE_PARALLEL=1 TP_SIZE=2 CP_SIZE=1 EP_SIZE=1 DP_SIZE=1 PP_SIZE=1 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=2 --no-python "$test_bin"
+elif [[ "$mode" == "tp-attention-smoke" ]]; then
+    TP_SIZE=2 CP_SIZE=1 EP_SIZE=1 DP_SIZE=1 PP_SIZE=1 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=2 --no-python "$test_bin"
+elif [[ "$mode" == "mtp-dynamic-smoke" ]]; then
+    WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 TP_SIZE=1 CP_SIZE=1 EP_SIZE=1 DP_SIZE=1 PP_SIZE=1 \
+        "$test_bin"
+elif [[ "$mode" == "mtp-dp-smoke" ]]; then
+    TP_SIZE=1 CP_SIZE=1 EP_SIZE=1 DP_SIZE=2 PP_SIZE=1 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=2 --no-python "$test_bin"
+elif [[ "$mode" == "mtp-tp-smoke" ]]; then
+    TP_SIZE=2 CP_SIZE=1 EP_SIZE=1 DP_SIZE=1 PP_SIZE=1 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=2 --no-python "$test_bin"
+elif [[ "$mode" == "mtp-tp-ep-smoke" ]]; then
+    QWEN36_TEST_MTP_MOE=1 QWEN36_TEST_MTP_EP=1 \
+    QWEN36_EP_A2A=1 QWEN36_EP_A2A_SHARDED=1 QWEN36_EP_A2A_PACKED=1 \
+    TP_SIZE=2 CP_SIZE=1 EP_SIZE=2 DP_SIZE=1 PP_SIZE=1 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=4 --no-python "$test_bin"
+elif [[ "$mode" == "cp-gdn-smoke" ]]; then
+    TP_SIZE=1 CP_SIZE=2 EP_SIZE=1 DP_SIZE=1 PP_SIZE=1 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=2 --no-python "$test_bin"
+elif [[ "$mode" == "cp-attention-smoke" ]]; then
+    cp_test_size="${QWEN36_TEST_CP_SIZE:-2}"
+    QWEN36_TEST_CP_RING="${QWEN36_TEST_CP_RING:-0}" \
+    QWEN36_CP_FULL_ATTENTION_KV_GATHER="$([[ "${QWEN36_TEST_CP_RING:-0}" == "1" ]] && echo 0 || echo 1)" \
+    QWEN36_CP_FULL_ATTENTION_RING="$([[ "${QWEN36_TEST_CP_RING:-0}" == "1" ]] && echo 1 || echo 0)" \
+    TP_SIZE=1 CP_SIZE="$cp_test_size" EP_SIZE=1 DP_SIZE=1 PP_SIZE=1 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node="$cp_test_size" --no-python "$test_bin"
+elif [[ "$mode" == "pp-cp-comm-smoke" ]]; then
+    TP_SIZE=1 CP_SIZE=2 EP_SIZE=1 DP_SIZE=1 PP_SIZE=2 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=4 --no-python "$test_bin"
+elif [[ "$mode" == "pp-train-smoke" ]]; then
+    pp_train_world="${PP_TRAIN_WORLD:-2}"
+    if [[ "$pp_train_world" -lt 2 ]]; then
+        echo "PP_TRAIN_WORLD must be >= 2" >&2
+        exit 1
+    fi
+    TP_SIZE=1 CP_SIZE=1 EP_SIZE=1 DP_SIZE=1 PP_SIZE="$pp_train_world" \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node="$pp_train_world" --no-python "$test_bin"
+elif [[ "$mode" == "tri-smoke" || "$mode" == "tri-replicated-smoke" ]]; then
+    sharded_a2a=1
+    if [[ "$mode" == "tri-replicated-smoke" ]]; then
+        sharded_a2a=0
+    fi
+    TP_SIZE=2 EP_SIZE=2 DP_SIZE=2 RUSTRAIN_DATA_PARALLEL=1 \
+    QWEN36_EP_A2A=1 QWEN36_EP_A2A_SHARDED="$sharded_a2a" \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=8 --no-python "$test_bin"
+elif [[ "$mode" == "smoke" || "$mode" == "gpu-metadata-smoke" || "$mode" == "bench" ]]; then
+    gpu_metadata_env=()
+    if [[ "$mode" == "gpu-metadata-smoke" ]]; then
+        gpu_metadata_env=(QWEN36_EP_A2A_GPU_METADATA=1)
+    fi
+    env "${gpu_metadata_env[@]}" TP_SIZE=2 EP_SIZE=2 DP_SIZE=1 \
+    QWEN36_EP_A2A=1 QWEN36_EP_A2A_SHARDED=1 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=4 --no-python "$test_bin"
+elif [[ "$mode" == "ep-bench" ]]; then
+    ep_bench_world="${EP_BENCH_WORLD:-4}"
+    TP_SIZE=1 EP_SIZE="$ep_bench_world" DP_SIZE=1 \
+    QWEN36_EP_A2A=1 QWEN36_EP_A2A_SHARDED=1 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node="$ep_bench_world" \
+            --no-python "$test_bin"
+else
+    TP_SIZE=1 EP_SIZE=2 DP_SIZE=1 \
+    QWEN36_EP_A2A=1 QWEN36_EP_A2A_SHARDED=1 \
+        "$python_bin" -m torch.distributed.run --standalone \
+            --nnodes=1 --nproc-per-node=2 --no-python "$test_bin"
+fi
diff --git a/src/main.rs b/src/main.rs
index 322a9333..abce26f9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -12,7 +12,7 @@ use rustrain_tch_tiny::tch_train;
 
 use std::path::{Path, PathBuf};
 
-use anyhow::{anyhow, bail, Context, Result};
+use anyhow::{Context, Result, anyhow, bail};
 use clap::{Parser, Subcommand};
 use tracing::info;
 
@@ -98,7 +98,7 @@ enum Command {
         #[arg(long, default_value_t = 1)]
         lora_rank: i32,
         #[arg(long, default_value_t = 60)]
-        duration: u64,  // seconds, 0 = single step
+        duration: u64, // seconds, 0 = single step
         #[arg(long, default_value = "512")]
         seq_len: usize,
     },
@@ -199,7 +199,14 @@ fn main() -> Result<()> {
             lora_rank,
             duration,
             seq_len,
-        } => run_ep_bench(metrics_dir, world_size, n_adapters, lora_rank, duration, seq_len),
+        } => run_ep_bench(
+            metrics_dir,
+            world_size,
+            n_adapters,
+            lora_rank,
+            duration,
+            seq_len,
+        ),
         Command::EpWorker {
             shm_name,
             rank,
@@ -231,8 +238,22 @@ fn dispatch_train(config_path: &Path, resume_from: Option<PathBuf>) -> Result<()
     validate_config(&config)?;
 
     let run_paths = prepare_run_directory(&config.run)?;
-    let _log_guard = init_logging(&run_paths.logs)?;
-    write_resolved_config(&config, &run_paths.resolved_config)?;
+    let world_size = std::env::var("WORLD_SIZE")
+        .ok()
+        .map(|value| value.parse::<usize>().context("WORLD_SIZE must be a usize"))
+        .transpose()?
+        .unwrap_or(1);
+    let rank = std::env::var("RANK")
+        .ok()
+        .map(|value| value.parse::<usize>().context("RANK must be a usize"))
+        .transpose()?
+        .unwrap_or(0);
+    let rank_log_dir =
+        rustrain_core::runtime::prepare_rank_log_directory(&run_paths, rank, world_size)?;
+    let _log_guard = init_logging(&rank_log_dir)?;
+    if rank == 0 {
+        write_resolved_config(&config, &run_paths.resolved_config)?;
+    }
 
     info!(config_path = %config_path.display(), "loaded config");
     info!(run_dir = %run_paths.root.display(), "created run directory");
@@ -595,7 +616,8 @@ fn dispatch_train(config_path: &Path, resume_from: Option<PathBuf>) -> Result<()
         let tp_size = config.parallel.tensor_model_parallel_size;
         let cp_size = config.parallel.context_parallel_size;
         if tp_size > 1 || cp_size > 1 {
-            let summary = rustrain_glm5::session_tp_cp::train_glm5_lora_sft_tp_cp_ep(&config, &run_paths)?;
+            let summary =
+                rustrain_glm5::session_tp_cp::train_glm5_lora_sft_tp_cp_ep(&config, &run_paths)?;
             println!("rustrain GLM-5.2 LoRA SFT TP+CP+EP complete");
             println!("run_dir: {}", run_paths.root.display());
             println!("adapter_checkpoint: {}", summary.adapter_output);
@@ -614,9 +636,9 @@ fn dispatch_train(config_path: &Path, resume_from: Option<PathBuf>) -> Result<()
         return Ok(());
     }
 
-    if is_tch && arch == "qwen3_6_lora_sft" {
+    if is_tch && matches!(arch, "qwen3_5_lora_sft" | "qwen3_6_lora_sft") {
         let summary = rustrain_qwen3_6::session::train_qwen3_6_lora_sft(&config, &run_paths)?;
-        println!("rustrain Qwen3.6 LoRA SFT complete");
+        println!("rustrain Qwen3.5/3.6 LoRA SFT complete");
         println!("run_dir: {}", run_paths.root.display());
         println!("adapter_checkpoint: {}", summary.adapter_output);
         println!("initial_loss: {:.9}", summary.initial_loss);
@@ -625,9 +647,9 @@ fn dispatch_train(config_path: &Path, resume_from: Option<PathBuf>) -> Result<()
         return Ok(());
     }
 
-    if is_tch && arch == "qwen3_6_lora_sft_ep" {
+    if is_tch && matches!(arch, "qwen3_5_lora_sft_ep" | "qwen3_6_lora_sft_ep") {
         let summary = rustrain_qwen3_6::session::train_qwen3_6_lora_sft_ep(&config, &run_paths)?;
-        println!("rustrain Qwen3.6 LoRA SFT EP complete");
+        println!("rustrain Qwen3.5/3.6 LoRA SFT EP complete");
         println!("run_dir: {}", run_paths.root.display());
         println!("adapter_checkpoint: {}", summary.adapter_output);
         println!("initial_loss: {:.9}", summary.initial_loss);
@@ -641,8 +663,8 @@ fn dispatch_train(config_path: &Path, resume_from: Option<PathBuf>) -> Result<()
 }
 
 fn run_server(http_port: u16, grpc_port: u16, metrics_dir: PathBuf) -> Result<()> {
-    use rustrain_server::{api, grpc, state::SessionManager};
     use rustrain_server::grpc::train::train_service_server::TrainServiceServer;
+    use rustrain_server::{api, grpc, state::SessionManager};
 
     std::fs::create_dir_all(&metrics_dir)?;
     let manager = std::sync::Arc::new(SessionManager::new(metrics_dir.clone()));
@@ -709,6 +731,7 @@ fn run_ep_server(
     rt.block_on(async move {
         let app_state = std::sync::Arc::new(api::EpAppState {
             coordinator: coordinator.clone(),
+            world_size,
         });
         let http_router = api::ep_router(app_state);
         let http_listener = tokio::net::TcpListener::bind(&http_addr)
@@ -740,8 +763,10 @@ fn run_ep_bench(
 
     std::fs::create_dir_all(&metrics_dir)?;
 
-    info!("EP bench: world_size={}, n_adapters={}, rank={}, duration={}s, seq={}",
-          world_size, n_adapters, lora_rank, duration, seq_len);
+    info!(
+        "EP bench: world_size={}, n_adapters={}, rank={}, duration={}s, seq={}",
+        world_size, n_adapters, lora_rank, duration, seq_len
+    );
 
     let coordinator = EpCoordinator::launch(world_size, metrics_dir.clone())
         .map_err(|e| anyhow!("Failed to launch EP workers: {}", e))?;
@@ -752,7 +777,8 @@ fn run_ep_bench(
         .unwrap_or_else(|_| "/mnt/workspace/huggingface/hub/models--Qwen--Qwen3.6-35B-A3B/snapshots/995ad96eacd98c81ed38be0c5b274b04031597b0".to_string());
     let lt = "[\"linear_attention\",\"linear_attention\",\"linear_attention\",\"full_attention\"]";
     let lt_full = std::iter::repeat(lt).take(8).collect::<Vec<_>>().join(",");
-    let config_toml = format!(r#"
+    let config_toml = format!(
+        r#"
 [run]
 name="bench"
 seed=42
@@ -806,10 +832,13 @@ target_modules=["q_proj","k_proj","v_proj","o_proj","in_proj_qkv","in_proj_z","o
 [data]
 kind="instruction_jsonl"
 paths=["/tmp/qwen3_6_test.jsonl"]
-"#);
+"#
+    );
 
     eprintln!("[bench] create session...");
-    let _ = coordinator.dispatch(&EpCommand::CreateSession { session_id: sid.to_string() });
+    let _ = coordinator.dispatch(&EpCommand::CreateSession {
+        session_id: sid.to_string(),
+    });
 
     eprintln!("[bench] load_model...");
     match coordinator.dispatch(&EpCommand::LoadModel {
@@ -817,7 +846,7 @@ paths=["/tmp/qwen3_6_test.jsonl"]
         model_path: model_path.to_string(),
         config_toml: config_toml.clone(),
     }) {
-        EpResult::Ok => {},
+        EpResult::Ok => {}
         EpResult::Error(e) => bail!("load_model failed: {}", e),
         _ => bail!("load_model unexpected result"),
     }
@@ -833,22 +862,31 @@ paths=["/tmp/qwen3_6_test.jsonl"]
     match coordinator.dispatch(&EpCommand::InitLora {
         session_id: sid.to_string(),
         rank: 8,
-        alpha: 16,
+        alpha: 16.0,
         target_layers: vec![],
-        target_modules: vec!["q_proj".to_string(),"k_proj".to_string(),"v_proj".to_string(),
-                             "o_proj".to_string(),"in_proj_qkv".to_string(),
-                             "in_proj_z".to_string(),"out_proj".to_string()],
+        target_modules: vec![
+            "q_proj".to_string(),
+            "k_proj".to_string(),
+            "v_proj".to_string(),
+            "o_proj".to_string(),
+            "in_proj_qkv".to_string(),
+            "in_proj_z".to_string(),
+            "out_proj".to_string(),
+        ],
         lr: 0.0001,
         beta1: 0.9,
         beta2: 0.999,
         eps: 0.00000001,
     }) {
-        EpResult::Count(_) => {},
+        EpResult::Count(_) => {}
         EpResult::Error(e) => bail!("init_lora failed: {}", e),
         _ => bail!("init_lora unexpected result"),
     }
 
-    eprintln!("[bench] batch_add_lora ({} adapters, rank={})...", n_adapters, lora_rank);
+    eprintln!(
+        "[bench] batch_add_lora ({} adapters, rank={})...",
+        n_adapters, lora_rank
+    );
     match coordinator.dispatch(&EpCommand::BatchAddLora {
         session_id: sid.to_string(),
         count: n_adapters,
@@ -856,16 +894,46 @@ paths=["/tmp/qwen3_6_test.jsonl"]
         alpha: (lora_rank * 2) as f64,
         target_layers: vec![],
         target_modules: "".to_string(),
+        optimizer_lr: None,
+        optimizer_beta1: None,
+        optimizer_beta2: None,
+        optimizer_eps: None,
     }) {
         EpResult::Count(n) => eprintln!("[bench] added {} adapters", n),
         EpResult::Error(e) => bail!("batch_add_lora failed: {}", e),
         _ => bail!("batch_add_lora unexpected result"),
     }
 
-    // Build input tensors
-    let ids: Vec<i64> = vec![1; seq_len];
-    let mask: Vec<i64> = { let mut m = vec![0i64; 20]; m.extend(vec![1i64; seq_len - 20]); m };
-    let attn: Vec<i64> = vec![1; seq_len];
+    // Build one n_total-row tenant batch for every independent source coordinate.
+    let topology = rustrain_parallel::topology::ParallelTopology::from_env_with_world_size(
+        world_size,
+    )?;
+    let ep_source_sharded = std::env::var("QWEN36_EP_A2A_SHARDED")
+        .map(|value| !value.is_empty() && value != "0")
+        .unwrap_or(false);
+    let source_parallel_size = if ep_source_sharded {
+        topology
+            .data_parallel_size()
+            .checked_mul(topology.expert_model_parallel_size())
+            .ok_or_else(|| anyhow!("source-parallel size overflowed usize"))?
+    } else {
+        topology.data_parallel_size()
+    };
+    let global_batch_size = usize::try_from(n_adapters)
+        .map_err(|_| anyhow!("n_adapters must be positive"))?
+        .checked_mul(source_parallel_size)
+        .ok_or_else(|| anyhow!("benchmark global batch size overflowed usize"))?;
+    if global_batch_size == 0 {
+        bail!("n_adapters and source-parallel size must be positive");
+    }
+    let tensor_elements = seq_len
+        .checked_mul(global_batch_size)
+        .ok_or_else(|| anyhow!("benchmark tensor element count overflowed usize"))?;
+    let ids = vec![1i64; tensor_elements];
+    let mut mask_row = vec![1i64; seq_len];
+    mask_row[..20.min(seq_len)].fill(0);
+    let mask = mask_row.repeat(global_batch_size);
+    let attn = vec![1i64; tensor_elements];
 
     // Warmup
     eprintln!("[bench] warmup...");
@@ -875,21 +943,31 @@ paths=["/tmp/qwen3_6_test.jsonl"]
         input_ids: ids.clone(),
         target_mask: mask.clone(),
         attention_mask: attn.clone(),
+        batch_size: global_batch_size,
         seq_len,
         n_total: n_adapters,
         lora_rank,
+        adapter_ids: vec![],
+        expected_steps: vec![],
     }) {
-        EpResult::Loss(l) => l,
-        EpResult::Error(e) => { bail!("warmup failed: {}", e); }
+        EpResult::Train { loss, .. } => loss,
+        EpResult::Error(e) => {
+            bail!("warmup failed: {}", e);
+        }
         _ => bail!("warmup unexpected result"),
     };
     let warmup_ms = t0.elapsed().as_millis();
-    eprintln!("[bench] warmup: loss={:.6} time={}ms", warmup_loss, warmup_ms);
+    eprintln!(
+        "[bench] warmup: loss={:.6} time={}ms",
+        warmup_loss, warmup_ms
+    );
 
     if duration == 0 {
         // Single step only
         eprintln!("[bench] single step done");
-        let _ = coordinator.dispatch(&EpCommand::DeleteSession { session_id: sid.to_string() });
+        let _ = coordinator.dispatch(&EpCommand::DeleteSession {
+            session_id: sid.to_string(),
+        });
         return Ok(());
     }
 
@@ -907,19 +985,29 @@ paths=["/tmp/qwen3_6_test.jsonl"]
             input_ids: ids.clone(),
             target_mask: mask.clone(),
             attention_mask: attn.clone(),
+            batch_size: global_batch_size,
             seq_len,
             n_total: n_adapters,
             lora_rank,
+            adapter_ids: vec![],
+            expected_steps: vec![],
         }) {
-            EpResult::Loss(l) => {
-                losses.push(l);
+            EpResult::Train { loss, .. } => {
+                losses.push(loss);
                 total_adapters += n_adapters as i64;
                 total_steps += 1;
                 let elapsed = start.elapsed().as_secs_f64();
                 if total_steps % 5 == 0 || total_steps == 1 {
                     let rate = total_adapters as f64 / elapsed;
-                    eprintln!("  step {}: loss={:.6} time={}ms total={} adp in {:.0}s = {:.1} adp/s",
-                              total_steps, l, t0.elapsed().as_millis(), total_adapters, elapsed, rate);
+                    eprintln!(
+                        "  step {}: loss={:.6} time={}ms total={} adp in {:.0}s = {:.1} adp/s",
+                        total_steps,
+                        loss,
+                        t0.elapsed().as_millis(),
+                        total_adapters,
+                        elapsed,
+                        rate
+                    );
                 }
             }
             EpResult::Error(e) => {
@@ -935,17 +1023,30 @@ paths=["/tmp/qwen3_6_test.jsonl"]
 
     eprintln!();
     eprintln!("============================================================");
-    eprintln!("  RESULTS: {} adapters, rank={}, EP={}", n_adapters, lora_rank, world_size);
+    eprintln!(
+        "  RESULTS: {} adapters, rank={}, EP={}",
+        n_adapters, lora_rank, world_size
+    );
     eprintln!("============================================================");
     eprintln!("  Duration: {:.0}s ({:.1} min)", elapsed, elapsed / 60.0);
     eprintln!("  Steps: {}", total_steps);
     eprintln!("  Total adapters processed: {}", total_adapters);
-    eprintln!("  Throughput: {:.2} adapters/s ({:.0} adapters/min)", rate, rate * 60.0);
+    eprintln!(
+        "  Throughput: {:.2} adapters/s ({:.0} adapters/min)",
+        rate,
+        rate * 60.0
+    );
     if !losses.is_empty() {
-        eprintln!("  Loss: {:.6} -> {:.6}", losses[0], losses[losses.len()-1]);
+        eprintln!(
+            "  Loss: {:.6} -> {:.6}",
+            losses[0],
+            losses[losses.len() - 1]
+        );
     }
     eprintln!("  Failures: {}", total_steps - losses.len() as i64);
 
-    let _ = coordinator.dispatch(&EpCommand::DeleteSession { session_id: sid.to_string() });
+    let _ = coordinator.dispatch(&EpCommand::DeleteSession {
+        session_id: sid.to_string(),
+    });
     Ok(())
 }