diff --git a/fastgen/methods/elastification/PLAN.md b/fastgen/methods/elastification/PLAN.md
new file mode 100644
index 0000000..d697ee9
--- /dev/null
+++ b/fastgen/methods/elastification/PLAN.md
@@ -0,0 +1,372 @@
+# Elastification for FastGen — Design Plan
+
+Bringing the **Elastification / Flextron** one-shot NAS + trained router compression method
+into FastGen, targeting Wan diffusion transformers.
+
+- Reference paper: <https://arxiv.org/pdf/2511.16664>
+- Reference implementation: `Megatron-LM/megatron/elastification/`
+
+## Design principle: mirror the Megatron reference
+
+Match Megatron-LM's `elastification/` module 1:1 wherever practical — attribute
+names, method names, `__init__` ordering, method-body structure, config field
+names. The intent is frictionless side-by-side comparison for debugging and
+code review.
+
+Only deviate when Megatron requires machinery FastGen doesn't have (pipeline /
+tensor parallelism, TransformerEngine parallel linears, `MegatronModule`,
+`get_args()` globals, `hybrid_layer_pattern`, ModelOpt distillation), or when a
+scope reduction from this plan applies (no Mamba/MoE/memory-mode/heterogeneous
+Mamba). Every deviation is called out in an inline comment; parallels are not.
+
+### Known Megatron quirks
+
+Reference behaviors that look like bugs. Fixed here when they'd break
+default configurations; noted for later verification against Megatron's
+actual training scripts.
+
+- **`FlextronRouter.mlp_forward` unbound `scale` — FIXED here, needs
+  Megatron-side verification.** When `linear_scaler_start` /
+  `linear_scaler_end` / `train_iters` are `None` (scaler disabled),
+  Megatron's `mlp_forward` still references a local `scale` variable that
+  was never assigned → `NameError` on the first forward. The only safe
+  config that avoids all `scale` references is
+  `normalize_router_logits=True` **and** `len(mlp_int_list) <= 1` (no MLP
+  elasticity). Every other config either requires the scaler to be
+  enabled or hits the NameError. Our port defaults `scale = 1.0` when
+  `self.scaler is None` so the `scale * ...` multiplications become
+  no-ops. **TODO:** double-check whether Megatron's actual training runs
+  always enable the scaler (via `--linear-scaler-start` / `--linear-scaler-end`
+  CLI args) — if so, the "bug" is dormant in their production configs and
+  our default-1.0 fix is a strict superset of the reference behavior.
+
+## 1. Locked-in decisions for this port
+
+| Area | Decision | Rationale |
+|---|---|---|
+| Target network | Wan | Product priority. Uses diffusers `WanTransformer3DModel`. |
+| Router axes | `emb_int_list`, `mlp_int_list`, optional `layer_ranking_list` | Matches Megatron's actually-implemented axes. |
+| Attention heads axis | **Dropped** | Megatron doesn't implement it either (`flex_budget_utils.py:41` says "head elasticity itself is no longer supported"). Not needed for parity, avoids the `timm.Attention` / diffusers `attn_processor` forward-override work. |
+| Head_dim axis | **Dropped** | Never varied in Megatron. |
+| Mamba / MoE hooks | **Dropped** | Wan has neither. |
+| Layout | Single directory `fastgen/methods/elastification/` | Method won't compose with other step-distillation methods (per user), so no shared-lib split needed. |
+| Phases | All three (joint / freeze-model / freeze-router) via config flags | Matches Megatron. Simple `param.requires_grad` toggles. |
+| Distillation | Included; **external frozen teacher** | Uses FastGen's `build_teacher()` infra; mirrors Megatron. Separate Wan checkpoint loaded frozen; +1 teacher forward per step. |
+| Checkpoint loading | Reuse FastGen's `pretrained_model_path` path | `load_student_weights_and_ema()` handles it. |
+| Head_dim divisibility constraint | **No hard enforcement** | Not needed for correctness. Optional soft warning if a user picks an `emb_int` that would produce a non-standard sub-net shape at materialization time. |
+| Materialization (weight cropping) | **v2, separate PR** | Not needed for algorithm correctness. Ships real inference perf savings. |
+
+## 2. Wan integration points
+
+FastGen's `Wan` (`fastgen/networks/Wan/network.py`) wraps diffusers'
+`WanTransformer3DModel` and monkey-patches several methods via
+`override_transformer_forward` (line 831-846). Relevant to us:
+
+- `block.forward = block_forward` — FastGen's block replacement, handles Wan2.1 T2V
+  modulation + optional sCM `norm_temb`. This is what our hooks attach around.
+- `transformer.forward = classify_forward` — FastGen's transformer-level forward.
+  Supports `skip_layers` (guidance-time layer skipping) — **note: this is NOT our
+  router-driven layer-skip**. Don't route through it; use our own block-level hook.
+- `transformer.classify_forward_block_forward` — the per-block iteration loop.
+
+**Hook attach order.** Our `apply_elasticity_to_wan(net, cfg)` runs during
+`ElastificationModel.build_model()`, which happens after `Wan.__init__` (which runs
+`override_transformer_forward`). So hooks bind to the FastGen-overridden `block_forward`.
+This ordering is automatic; do not move elasticity setup earlier.
+
+**Module targets inside each `WanTransformerBlock`:**
+
+- `attn1` — self-attention (`WanAttention` with separate `to_q`, `to_k`, `to_v`,
+  `to_out.0`, `norm_q`, `norm_k`).
+- `attn2` — cross-attention over text embeddings. Same shape as `attn1` but K, V come from
+  the text encoder side.
+- `ffn` — diffusers `FeedForward` (need to check its internal structure at implementation
+  time — likely has `.net.0.proj` and `.net.2`).
+- `scale_shift_table` — AdaLN modulation tensor `[6, inner_dim]`; needs emb masking on the
+  `inner_dim` axis so shift/scale/gate operate on the pruned channels only.
+
+**Cross-attention specifics.** Router axes affect `attn2` as follows:
+
+- `emb` — mask Q input (video-side); mask `to_k` and `to_v` **output** rows to the emb
+  prefix (so K/V dims line up with Q's shrunk head layout); leave `encoder_hidden_states`
+  (text side) untouched — that's the frozen text encoder's output dim, not ours to prune.
+- `mlp_hidden` — no effect on cross-attn.
+- `layer_skip` — dropping a block skips its cross-attn too.
+
+## 3. File structure
+
+### Implementation (new)
+
+```
+fastgen/methods/elastification/
+├── __init__.py                       # exports ElastificationModel
+├── PLAN.md                           # this file
+├── README.md                         # end-user doc: search space, phases, inference
+├── config.py                         # FlextronConfig (attrs dataclass) + validation
+├── router.py                         # FlextronRouter (Gumbel-softmax, DP-seeded, tau decay)
+├── budget_utils.py                   # analytical Wan param count on soft router outputs
+├── hooks/
+│   ├── __init__.py                   # apply_elasticity_to_wan(net, cfg)
+│   ├── self_attention.py             # WanAttention as attn1
+│   ├── cross_attention.py            # WanAttention as attn2 (video-side)
+│   ├── ffn.py                        # emb + mlp-hidden mask
+│   ├── block.py                      # WanTransformerBlock: layer-skip + scale_shift_table
+│   └── stack.py                      # WanTransformer3DModel: norm_out + proj_out emb mask
+├── manager.py                        # ElastificationManager
+├── loss.py                           # combine_losses, self_distill helper
+├── model.py                          # ElastificationModel(FastGenModel)
+└── materialize.py                    # v2 (later PR): crop weights to fixed budget
+```
+
+### Config (new)
+
+```
+fastgen/configs/methods/
+└── config_elastification.py          # ModelConfig(BaseModelConfig) + Config(BaseConfig)
+```
+
+### Registration (edit existing)
+
+```
+fastgen/methods/__init__.py           # + from fastgen.methods.elastification.model import ElastificationModel
+```
+
+### Inference script (new)
+
+```
+scripts/inference/
+└── run_elastification_inference.py   # standalone: load ckpt, set override_selected_budget, sample
+```
+
+### Tests (new)
+
+```
+tests/elastification/
+├── __init__.py
+├── test_router.py                    # ~200 L
+├── test_budget_utils.py              # ~150 L
+├── test_apply_elasticity_to_wan.py   # ~200 L
+├── test_self_attention_hook.py       # ~150 L
+├── test_cross_attention_hook.py      # ~150 L
+├── test_ffn_hook.py                  # ~150 L
+├── test_block_hook.py                # ~150 L
+├── test_manager.py                   # ~150 L
+├── test_config.py                    # ~100 L
+├── test_model.py                     # ~200 L
+├── test_loss.py                      # ~100 L
+└── test_numerical_invariants.py      # ~300 L (checks Megatron doesn't have)
+```
+
+### Rough sizing
+
+- Implementation: 13 files, ~2500 lines
+- Config: 1 file, ~50 lines
+- Inference script: 1 file, ~100 lines
+- Tests: 13 files, ~2100 lines
+- **Total: ~28 new files, ~4750 lines**
+
+Comparable to Megatron's `elastification/` (~4800 impl + ~2300 tests), scoped down for
+the smaller search space (no Mamba, no MoE, no per-block heterogeneity, no TP router).
+
+## 4. Per-file design notes
+
+### `config.py`
+
+```python
+@attrs.define(slots=False)
+class FlextronConfig:
+    # ── Search space ──────────────────────────────────
+    emb_int_list: List[int] = []                  # e.g. [5120, 3840, 2560, 1280]
+    mlp_int_list: List[int] = []                  # e.g. [13824, 10368, 6912, 3456]
+    layer_ranking_list: Optional[List[int]] = None  # ordered block indices eligible to skip
+
+    # ── Budget ────────────────────────────────────────
+    budget_list: List[float] = [1.0]
+    budget_probs: Optional[List[float]] = None
+    loss_alpha: float = 1.0
+    router_beta: float = 1.0
+
+    # ── Router ────────────────────────────────────────
+    enable_router: bool = True
+    router_inter_dim: int = 128
+    tau_init: float = 1.0
+    tau_decay: float = 0.9999
+    hard_sample_th: float = 0.996
+    router_std: float = 0.1
+    lr_mult_router: float = 1.0
+    normalize_router_logits: bool = False
+    linear_scaler_start: Optional[float] = None
+    linear_scaler_end: Optional[float] = None
+
+    # ── Training schedule ─────────────────────────────
+    original_model_sample_prob: float = 0.33      # identity anchor
+
+    # ── Phase flags ───────────────────────────────────
+    freeze_model: bool = False
+    freeze_router: bool = False
+
+    # ── Distillation ──────────────────────────────────
+    distillation: bool = False                    # enable external-teacher distillation
+    distill_coeff: float = 1.0
+
+    # ── Eval / inference override ─────────────────────
+    override_selected_budget: Optional[float] = None
+    is_eval: bool = False
+
+    def __attrs_post_init__(self):
+        # Sort budget_list descending (Megatron invariant)
+        if self.budget_list:
+            order = sorted(range(len(self.budget_list)), key=lambda i: -self.budget_list[i])
+            self.budget_list = [self.budget_list[i] for i in order]
+            if self.budget_probs:
+                self.budget_probs = [self.budget_probs[i] for i in order]
+```
+
+### `router.py`
+
+`FlextronRouter(nn.Module)`. Same shape as Megatron's `FlextronRouter` stripped of tensor-parallel
+linear layers (FastGen uses DDP/FSDP, not TP). Just two `nn.Linear` layers per axis. Preserves:
+
+- Gumbel-softmax with `tau = tau_init * tau_decay^iteration`
+- DP-seeded Gumbel noise: seed is `base_seed + (dp_rank + fwd_count * dp_size) % router_gbs + iteration * 1000`, so each (rank, micro-step, iteration) triple draws distinct noise. Different DP ranks and different grad-accum micro-steps sample *different* sub-architectures at the same target budget — one training step covers `dp_size × grad_accum_rounds` distinct samples of the discrete choice space, increasing gradient-signal diversity per update. Deterministic across reruns.
+- Budget one-hot lookup + linear interpolation for intermediate budgets
+- `forward(budget) → dict[str, (soft_probs, chosen_int)]` per axis
+- 3 axes: `mlp`, `emb`, `layer_skip`
+
+### `budget_utils.py`
+
+Wan-specific analytical param count. Modeled on Megatron's
+`flex_budget_utils.py::get_num_parameters`. Accepts either concrete ints or router soft
+probs (differentiable). Formula per block:
+
+```
+attn1_params = 4 * hidden_size * hidden_size                            # to_q + to_k + to_v + to_out.0
+attn2_params = 2 * hidden_size * hidden_size                            # to_q + to_out.0
+              + 2 * text_encoder_dim * hidden_size                       # to_k + to_v (text side untouched)
+ffn_params   = 2 * hidden_size * mlp_hidden                              # up_proj + down_proj
+misc_params  = 6 * hidden_size + 3 * hidden_size                         # scale_shift_table + norms
+block_params = attn1_params + attn2_params + ffn_params + misc_params
+
+total_blocks = Σ (1 - skip_prob_i) * block_params
+```
+
+Params scale linearly with `hidden_size` and `mlp_hidden`. `num_heads` and `head_dim` don't
+appear because they never vary in the search space. Text-encoder dim is a constant.
+
+### `hooks/` — five hook managers
+
+Each follows the Megatron pattern: constructor stores config, `attach_hooks(module)`
+registers pre/post hooks, `set_elasticity_params(...)` receives router choices per step,
+`detach_hooks()` restores.
+
+- **`self_attention.py`** (attn1) — emb-mask on input, eps rescale on `norm_q`/`norm_k`,
+  emb mask + STE prob scaling on `to_out.0` output.
+- **`cross_attention.py`** (attn2) — like self-attention but Q input only (video side);
+  `to_k`/`to_v` output rows masked; text-side untouched.
+- **`ffn.py`** — emb mask input, mlp-hidden mask on intermediate (post-hook on first linear),
+  emb mask on final output.
+- **`block.py`** — layer-skip: pre-hook stashes input, post-hook returns stashed input if
+  the router selected this block for skipping. Also masks `scale_shift_table` on `inner_dim`.
+- **`stack.py`** — top-level `norm_out` eps rescale + `sqrt(k/N)` restore. `proj_out`
+  input is already masked upstream.
+
+`hooks/__init__.py::apply_elasticity_to_wan(net, cfg)` walks
+`net.transformer.blocks` and attaches every manager. Modeled directly on Megatron's
+`apply_flextron_elasticity_to_model`. Returns list of managers for the orchestrator.
+
+### `manager.py` — `ElastificationManager`
+
+Orchestrator. Modeled on Megatron's `FlextronModelManager`.
+
+- `__init__(net, cfg)`:
+  - Store net + cfg
+  - Build `FlextronRouter`
+  - Call `apply_elasticity_to_wan(net, cfg)` → store hook managers
+  - Pre-compute full-model param count for budget-loss normalization
+  - Optional: warn if any `emb_int_list` entry is not a multiple of `head_dim` (soft warning; not enforced)
+- `process_budget(budget) → soft_choices` — calls router forward
+- `push_to_hooks(soft_choices)` — forwards to every hook manager's `set_elasticity_params`
+- `compute_budget_loss(soft_choices, budget) → tensor` — delegates to `budget_utils`
+
+### `loss.py`
+
+- `budget_loss_func(current_soft_params, target_budget, full_params, budget_val)` —
+  `|current / (target * full) - 1|`, clipped inside 5% dead-zone, MSE anchor at `budget = 1.0`.
+- `distillation_loss_func(student_pred, teacher_pred)` — MSE.
+- `combine_losses(dsm_loss, budget_loss, distill_loss, alpha, distill_coeff) → total_loss, loss_map`.
+
+### `model.py` — `ElastificationModel(FastGenModel)`
+
+A FastGen method (subclass of FastGenModel) with a Megatron-style hook orchestrator inside.
+
+- `build_model()`:
+  1. `super().build_model()` builds `self.net` (Wan)
+  2. `self.load_student_weights_and_ema()` loads pretrained ckpt
+  3. `self.manager = ElastificationManager(self.net, self.config.elast)` — sets up router + hooks
+  4. If `distillation`: `self.build_teacher()` (FastGen infra)
+  5. Apply freeze flags: `self.net.requires_grad_(False)` if `freeze_model`;
+     `self.manager.router.requires_grad_(False)` if `freeze_router`
+- `single_train_step(data, iteration)`:
+  1. Prepare data (SFT-style)
+  2. Sample budget via DP-seeded choice (mirror Megatron's `get_grad_acc_based_random_choice`)
+  3. `soft_choices = self.manager.process_budget(budget)`
+  4. `budget_loss = self.manager.compute_budget_loss(soft_choices, budget)`
+  5. `self.manager.push_to_hooks(soft_choices)`
+  6. `net_pred = self.net(noisy, t, condition=cond_train)` — hooks fire inside
+  7. `dsm_loss = denoising_score_matching_loss(...)`
+  8. Distillation branch (if enabled)
+  9. Combine: `total_loss = dsm_loss + loss_alpha * budget_loss + distill_coeff * distill_loss`
+  10. Return `loss_map = {total_loss, dsm_loss, budget_loss, distill_loss, sampled_budget}`, `outputs`
+- `init_optimizers()`: parent builds `net_optimizer`; add `router_optimizer` with
+  `LR × lr_mult_router`.
+- `get_optimizers(iteration) → [net_optimizer, router_optimizer]` (respects freeze flags).
+- `optimizer_dict`, `scheduler_dict`, `model_dict`: include router entries for checkpointing.
+
+### `materialize.py` (v2, later PR)
+
+Post-training utility. Given trained elastic model + fixed budget, physically crop weights:
+
+- Freeze router at hard choice for the budget
+- Slice `attn1.to_q/to_k/to_v/to_out.0`, `attn2.to_q/to_out.0/norm_q` weights on emb axis;
+  slice `attn2.to_k/to_v` on output rows; slice `ffn.net.0.proj/net.2` on mlp-hidden axis;
+  slice `scale_shift_table` on `inner_dim`
+- Delete blocks whose skip prob > 0.5
+- Return a new `Wan` instance with cropped weights
+
+Delivers real inference perf savings (activation masking alone doesn't). Separate PR.
+
+## 5. Inference
+
+### At runtime
+
+Wan inference goes through `self.net.sample(noise, condition=..., ...)`, which internally
+calls `self.net.forward(noisy, t, condition)` per denoising step. Since hooks are attached
+to modules inside `self.net.transformer`, they fire on every forward automatically.
+
+Recipe:
+
+1. Set `config.model.elast.override_selected_budget = <b>`
+2. Set `config.model.elast.is_eval = True`
+3. Load elastic checkpoint (network + router)
+4. Call standard FastGen inference script
+5. Hooks apply the sub-net masking on every denoising step
+
+### Standalone entry point
+
+`scripts/inference/run_elastification_inference.py`:
+
+1. Parse args + FlextronConfig
+2. Instantiate `ElastificationModel` from config
+3. Load ckpt
+4. Pin `override_selected_budget` and `is_eval`
+5. Call `model.net.sample(noise, condition=..., ...)`
+
+Megatron ships no equivalent script (the `--is-flex-eval` flag exists but no in-tree
+caller uses it) — our port ships one from the start.
+
+### Caveat: no wall-clock savings from masking alone
+
+Activation masking does not reduce compute. Attention still runs at full head count,
+FFN still projects to full `mlp_hidden`. For real perf savings, use v2 materialization
+after training.
+
diff --git a/fastgen/methods/elastification/budget_utils.py b/fastgen/methods/elastification/budget_utils.py
new file mode 100644
index 0000000..a2bae03
--- /dev/null
+++ b/fastgen/methods/elastification/budget_utils.py
@@ -0,0 +1,136 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Analytical parameter count for Wan sub-networks under Elastification.
+
+Ports Megatron-LM's `flex_budget_utils.py::get_num_parameters`.
+
+Router-controlled widths (`hidden_size`, `ffn_hidden_size`) may be either
+plain ints (for a concrete sub-net) or `torch.Tensor` (for the soft expected
+value produced by `soft_probs @ candidate_list_tensor`, giving a
+differentiable budget-loss term).
+
+Wan-scope drops:
+  - No Mamba, no MoE → drops `mamba_*`, `moe_*` inputs, `flex_hetero_mamba`
+    / `flex_hetero_moe_expert` branches, `moe_active` vs `moe_all` distinction,
+    and returns a single count (not the `(all, active)` tuple).
+  - No hybrid layer-pattern string → iterates `range(num_layers)` directly.
+  - No GQA → drops `kv_channels` + `num_query_groups` split; uses `head_dim`
+    such that `num_attention_heads * head_dim = hidden_size` at full width.
+
+Wan additions (no Megatron analog):
+  - `text_encoder_dim` — attn2 cross-attention K/V input width.
+  - `in_channels` / `out_channels` / `patch_dim` — patch-embed + proj-out.
+  - `time_freq_dim` — sinusoidal-embedding intermediate.
+  - Cross-attention block per WanTransformerBlock.
+  - AdaLN modulation table `[6, hidden]` per block.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Union
+
+import torch
+
+
+ScalarOrTensor = Union[int, float, torch.Tensor]
+
+
+def get_num_parameters(
+    num_layers: int = 0,
+    num_attention_heads: int = 0,
+    head_dim: int = 0,
+    hidden_size: ScalarOrTensor = 0,
+    ffn_hidden_size: ScalarOrTensor = 0,
+    # ── Wan additions ─────────────────────────────────────────────────
+    text_encoder_dim: int = 0,
+    in_channels: int = 0,
+    out_channels: int = 0,
+    patch_dim: int = 0,
+    time_freq_dim: int = 256,
+    layer_skip_probs: Optional[torch.Tensor] = None,
+) -> ScalarOrTensor:
+
+    norm_multiplier = 1
+
+    # Wan input / output boilerplate. Megatron LLMs have `embedding` +
+    # `final_layernorm` + `output_layer` here; Wan has patch-embed on the
+    # input side and a projection back to patch space on the output side,
+    # plus condition-embedder projections that scale with hidden_size.
+    patch_embed = in_channels * patch_dim * hidden_size
+    time_embed = time_freq_dim * hidden_size + hidden_size * hidden_size
+    text_embed = text_encoder_dim * hidden_size + hidden_size * hidden_size
+    final_layernorm = hidden_size * 1
+    proj_out = hidden_size * (patch_dim * out_channels)
+
+    if isinstance(ffn_hidden_size, int):
+        flex_hetero_ffn = False
+    else:
+        flex_hetero_ffn = ffn_hidden_size.shape[0] != 1
+
+    # FFN (mirrors Megatron's `moe_all`/`moe_active`, collapsed to a single
+    # value because Wan has no MoE).
+
+    if flex_hetero_ffn:
+        ffn_all = []
+        for i in range(ffn_hidden_size.shape[0]):
+            pre_mlp_ln = norm_multiplier * hidden_size
+            linear_fc1 = ffn_hidden_size[i] * hidden_size
+            linear_fc2 = ffn_hidden_size[i] * hidden_size
+            ffn_all.append(pre_mlp_ln + linear_fc1 + linear_fc2)
+    else:
+        pre_mlp_ln = norm_multiplier * hidden_size
+        linear_fc1 = ffn_hidden_size * hidden_size
+        linear_fc2 = ffn_hidden_size * hidden_size
+        ffn_all = pre_mlp_ln + linear_fc1 + linear_fc2
+
+    # ATT (self-attention — `attn1`). Wan uses full attention with
+    # `num_attention_heads * head_dim = hidden_size`, so we compute
+    # `linear_qkv = 3 * num_attention_heads * head_dim * hidden_size` directly
+    # rather than the `(num_heads + 2*num_query_groups) * kv_channels * hidden`
+    # GQA-aware formula Megatron uses.
+
+    input_ln = norm_multiplier * hidden_size
+    linear_proj = num_attention_heads * head_dim * hidden_size
+    linear_qkv = 3 * num_attention_heads * head_dim * hidden_size
+    att = input_ln + linear_proj + linear_qkv
+
+    # CROSS-ATT (`attn2`) — new relative to Megatron.
+    # `linear_q` and `linear_out` both project on the video-side hidden;
+    # `linear_k` and `linear_v` project from the frozen text-encoder output.
+
+    input_ln_cross = norm_multiplier * hidden_size
+    linear_q_cross = num_attention_heads * head_dim * hidden_size
+    linear_proj_cross = num_attention_heads * head_dim * hidden_size
+    linear_kv_cross = 2 * text_encoder_dim * hidden_size
+    cross = input_ln_cross + linear_q_cross + linear_proj_cross + linear_kv_cross
+
+    # AdaLN modulation `scale_shift_table` `[6, hidden]` per block.
+    adaln = 6 * hidden_size
+
+    # Per-block loop mirrors Megatron's `for i, c in enumerate(hybrid_pattern):`
+    # but Wan has no pattern — every block is the same type. Optional
+    # `layer_skip_probs[i]` weights each block's contribution by
+    # `(1 - skip_prob)` (the "phantom" skip axis).
+
+    all_params = 0
+    for i in range(num_layers):
+        if flex_hetero_ffn:
+            block = att + cross + adaln + ffn_all[i]
+        else:
+            block = att + cross + adaln + ffn_all
+
+        if layer_skip_probs is not None:
+            block = block * (1 - layer_skip_probs[i])
+
+        all_params += block
+
+    return (
+        patch_embed
+        + time_embed
+        + text_embed
+        + all_params
+        + final_layernorm
+        + proj_out
+    )
diff --git a/fastgen/methods/elastification/config.py b/fastgen/methods/elastification/config.py
new file mode 100644
index 0000000..389fd6c
--- /dev/null
+++ b/fastgen/methods/elastification/config.py
@@ -0,0 +1,138 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+FlextronConfig — all Elastification config fields in one place.
+
+Ports `FlextronConfig` from `megatron/elastification/flextron_config.py`.
+
+Section order, field order, field names, and defaults mirror the reference.
+Dropped (out of Wan scope): all Mamba fields, all MoE fields, all memory-mode
+fields, `basemodel_type`. Added (Wan needs): `lr_mult_router`, `train_iters`,
+`num_mlp_blocks` — noted inline.
+
+Framework: `@attrs.define(slots=False)` is used instead of Megatron's
+`@dataclass` so FlextronConfig composes with FastGen's attrs-based ModelConfig
+hierarchy (see `fastgen/configs/methods/config_sft.py`).
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+import attrs
+
+
+@attrs.define(slots=False)
+class FlextronConfig:
+    # ── Core flags ────────────────────────────────────────────────────────────
+    flextron: bool = False
+    binary_mask: bool = False
+    add_skipping: bool = False
+    no_attn_skip: bool = False
+    slice: bool = False
+    soft_mask: bool = False
+
+    # ── Router ────────────────────────────────────────────────────────────────
+    enable_router: bool = False
+    router_inter_dim: int = 128
+    hard_sample_th: float = 0.996
+    router_beta: float = 1.0
+    loss_alpha: float = 1.0
+    tau_init: float = 1.0
+    tau_decay: float = 0.9999
+    router_std: float = 0.1
+    router_gbs: int = 32
+    normalize_router_logits: bool = False
+    linear_scaler_start: Optional[float] = None
+    linear_scaler_end: Optional[float] = None
+    # Megatron: CLI arg on the training script.
+    lr_mult_router: float = 1.0
+    # Megatron: pulled from `args.train_iters`.
+    train_iters: Optional[int] = None
+
+    # ── Budget ────────────────────────────────────────────────────────────────
+    budget_probs: Optional[List[float]] = None
+    budget_list: Optional[List[float]] = None
+    budget_type: str = 'param'
+    disable_budget: bool = False
+
+    # ── Training / eval control ───────────────────────────────────────────────
+    is_flex_eval: bool = False
+    freeze_router: bool = False
+    freeze_model: bool = False
+    curr_iteration: Optional[int] = None
+    original_model_sample_prob: float = 0.33
+    override_selected_budget: Optional[List[float]] = None
+
+    # ── Layer-skip constraints ────────────────────────────────────────────────
+    skip_num_attn_layer_constraint: Optional[int] = None
+    skip_total_layer_constraint: Optional[int] = None
+    layer_ranking_list: Optional[List[int]] = None
+    # Megatron inherits `num_layers` from TransformerConfig; the manager
+    # populates this from `net.transformer.blocks`.
+    num_layers: Optional[int] = None
+
+    # ── Force overrides (eval / frozen-router mode) ───────────────────────────
+    force_router_skip: Optional[List[int]] = None
+    force_mlp: Optional[List[float]] = None
+    force_emb: Optional[List[float]] = None
+
+    # ── Choice lists (converted to int at model-setup time) ───────────────────
+    mlp_per_list: Optional[List[float]] = None
+    emb_per_list: Optional[List[float]] = None
+    mlp_int_list: Optional[List[int]] = None
+    emb_int_list: Optional[List[int]] = None
+
+    # ── Heterogeneous per-layer routing ───────────────────────────────────────
+    flex_hetero_ffn: bool = False
+    # Megatron derives per-block counts from `hybrid_layer_pattern`; Wan has no
+    # such string, so the manager populates this from `net.transformer.blocks`.
+    num_mlp_blocks: Optional[int] = None
+
+    # ── Distillation ──────────────────────────────────────────────────────────
+    distillation: bool = False
+    distill_coeff: float = 0.0
+    distill_only: bool = False
+
+    def __attrs_post_init__(self):
+        # Mirrors Megatron's `validate_flextron_per_int_lists` +
+        # `sort_budget_list_descending` from `elastification/arguments.py`,
+        # inlined here because we don't have a separate args-injection step.
+
+        # Per-list / int-list mutual exclusion. Defaults per-list to [1.0]
+        # (elasticity off) when neither is set on an axis.
+        pairs = (
+            ('mlp', 'mlp_per_list', 'mlp_int_list'),
+            ('emb', 'emb_per_list', 'emb_int_list'),
+        )
+        for name, per_attr, int_attr in pairs:
+            per_val = getattr(self, per_attr)
+            int_val = getattr(self, int_attr)
+            per_set = per_val is not None
+            int_set = int_val is not None
+            if per_set:
+                for x in per_val:
+                    assert 0.0 <= x <= 1.0, (
+                        f'{per_attr} values must be in [0, 1], got {x}.'
+                    )
+            assert not (per_set and int_set), (
+                f'Use either {per_attr} or {int_attr} for {name}, not both.'
+            )
+            if not per_set and not int_set:
+                setattr(self, per_attr, [1.0])
+
+        # Sort budget_list descending; permute budget_probs to match.
+        if self.budget_list is not None and len(self.budget_list) > 1:
+            if self.budget_probs is not None:
+                assert len(self.budget_probs) == len(self.budget_list), (
+                    f'budget_probs length {len(self.budget_probs)} does not '
+                    f'match budget_list length {len(self.budget_list)}'
+                )
+            order = sorted(
+                range(len(self.budget_list)),
+                key=lambda i: -self.budget_list[i],
+            )
+            self.budget_list = [self.budget_list[i] for i in order]
+            if self.budget_probs is not None:
+                self.budget_probs = [self.budget_probs[i] for i in order]
diff --git a/fastgen/methods/elastification/router.py b/fastgen/methods/elastification/router.py
new file mode 100644
index 0000000..0ae948f
--- /dev/null
+++ b/fastgen/methods/elastification/router.py
@@ -0,0 +1,400 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Elastic router for the Elastification method.
+
+Ports Megatron-LM's FlextronRouter
+(megatron/elastification/router/hybrid_flex_router.py) to FastGen, adapted for
+Wan diffusion transformers. Compared to the reference implementation:
+
+- Uses plain nn.Linear instead of TransformerEngine parallel linears (FastGen
+  is DDP/FSDP-only; no tensor parallelism).
+- Only three axes: mlp, emb, skip (Wan has no Mamba or MoE).
+- Layer-skip axis is retained to match the reference exactly: the router
+  produces a router_skip output that feeds the analytical budget loss, but
+  nothing wires it into an actual forward-pass skip. It is a phantom axis
+  whose gradient flows through the budget-loss param count only.
+- No per-block heterogeneity ("flex_hetero_*" branches dropped).
+- No pipeline-parallel weight sync.
+- Iteration is passed explicitly to forward() rather than pulled from a global
+  args object.
+"""
+
+from __future__ import annotations
+
+import random
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from fastgen.methods.elastification.config import FlextronConfig
+from fastgen.utils.distributed import get_rank, world_size
+
+
+class FlextronRouter(nn.Module):
+    """Router that maps a scalar budget to per-axis Gumbel-softmax choices.
+
+    Structure:
+      - self.gate_emb    : Sequential producing logits over emb_int_list
+      - self.gate_mlp    : Sequential producing logits over mlp_int_list
+      - self.gate_skip_layer (optional) : Sequential producing skip logits
+
+    Sampling is DP-seeded: each (dp_rank, fwd_pass_count, iteration) triple
+    draws a distinct Gumbel noise pattern via a deterministic seed. Different
+    DP ranks and different grad-accum micro-steps therefore sample *different*
+    sub-architectures around the same target budget — over one training step
+    the router sees `dp_size * grad_accum_rounds` distinct sub-arch samples,
+    which increases gradient-signal diversity per update.
+    """
+
+    def __init__(
+        self,
+        config: FlextronConfig,
+        gumbel_base_seed: int = 42,
+    ):
+        super().__init__()
+        self.config = config
+
+        self.input_dim = len(self.config.budget_list)
+        self.n_dim = self.config.router_inter_dim
+        self.budget_map = {
+            item: torch.tensor(idx)
+            for idx, item in enumerate(self.config.budget_list)
+        }
+
+        # Initialize DP-aware Gumbel softmax
+        self._init_dp_gumbel_softmax(base_seed=gumbel_base_seed)
+
+        # Create init method for router layers
+        # `init_method_normal` is inlined locally rather than importing from
+        # `megatron.core.utils`.
+        self.init_method = self._init_method_normal(self.config.router_std)
+
+        self.add_router_for_mlp()
+        self.add_router_for_emb()
+        if self.config.add_skipping:
+            self.add_router_for_skipping()
+
+        self.hard_sample_th = self.config.hard_sample_th
+
+        self.add_scaler_schedule()
+
+        # FastGen has no PP/TP so `world_size()` is the DP world size.
+        self.dp_size = world_size()
+        self.fwd_pass_count = 0
+
+    # ── Helpers ────────────────────────────────────────────────────────
+
+    def _init_dp_gumbel_softmax(self, base_seed: int = 42) -> None:
+        """Initialize DP-aware Gumbel-softmax state.
+
+        FastGen has no PP/TP so `get_rank()` is the DP rank directly.
+        """
+        self.dp_rank = get_rank()
+        self.gumbel_base_seed = base_seed
+
+    def _init_method_normal(self, std: float):
+        """Local inline replacement for `megatron.core.utils.init_method_normal`."""
+        def init_fn(tensor: torch.Tensor) -> torch.Tensor:
+            return nn.init.normal_(tensor, mean=0.0, std=std)
+        return init_fn
+
+    def _dp_gumbel_softmax(self, logits, tau=1.0, hard=False, curr_iteration=0):
+        """DP-aware Gumbel softmax that uses different random seeds per DP rank and iteration"""
+        # Create unique seed for this iteration and DP rank
+
+        seed = (
+            self.gumbel_base_seed
+            + (self.dp_rank + self.fwd_pass_count * self.dp_size) % self.config.router_gbs
+            + curr_iteration * 1000
+        )
+        # torch.manual_seed seeds both CPU and CUDA generators globally, so we
+        # must save/restore both - otherwise the CUDA RNG leaks the deterministic
+        # state we set here into other CUDA random ops elsewhere in the model.
+        cpu_state = torch.get_rng_state()
+        cuda_state = (
+            torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
+        )
+        torch.manual_seed(seed)
+
+        try:
+            return F.gumbel_softmax(logits, tau=tau, hard=hard)
+        finally:
+            torch.set_rng_state(cpu_state)
+            if cuda_state is not None:
+                torch.cuda.set_rng_state_all(cuda_state)
+
+    def _create_linear_layer(
+        self, input_size: int, output_size: int, bias: bool = False
+    ) -> nn.Linear:
+        """Standard router linear.
+
+        Strips the TransformerEngine parallel-linear branch and the
+        `is_first_layer` distinction from Megatron's version (both irrelevant
+        without tensor parallelism).
+        """
+        layer = nn.Linear(input_size, output_size, bias=bias)
+        self.init_method(layer.weight)
+        return layer
+
+    def add_router_for_mlp(self) -> None:
+        mlp_list = self.config.mlp_int_list
+        if self.config.flex_hetero_ffn:
+            # Megatron reads block count from `hybrid_layer_pattern`; Wan uses
+            # `num_mlp_blocks` (populated by the manager from the model).
+            num_mlp = self.config.num_mlp_blocks
+            gate_mlp_layer_list = [
+                self._create_linear_layer(self.input_dim, self.n_dim, bias=False),
+                nn.LeakyReLU(0.1),
+                self._create_linear_layer(
+                    self.n_dim, len(mlp_list) * num_mlp, bias=False
+                ),
+            ]
+            # Defensive last-layer bias init (dead with plain nn.Linear + bias=False;
+            # kept for parity with Megatron in case a caller flips bias on).
+            if (
+                hasattr(gate_mlp_layer_list[-1], 'bias')
+                and gate_mlp_layer_list[-1].bias is not None
+            ):
+                last_layer_bias = [0.00 for _ in range(len(mlp_list))]
+                last_layer_bias[-1] = 1.00
+                gate_mlp_layer_list[-1].bias.data = torch.tensor(
+                    last_layer_bias,
+                    dtype=gate_mlp_layer_list[-1].weight.dtype,
+                    device=gate_mlp_layer_list[-1].weight.device,
+                ).repeat(num_mlp)
+        else:
+            gate_mlp_layer_list = [
+                self._create_linear_layer(self.input_dim, self.n_dim, bias=False),
+                nn.LeakyReLU(0.1),
+                self._create_linear_layer(self.n_dim, len(mlp_list), bias=False),
+            ]
+        self.gate_mlp = nn.Sequential(*gate_mlp_layer_list)
+
+    def add_router_for_emb(self) -> None:
+        emb_list = self.config.emb_int_list
+        gate_emb_layer_list = [
+            self._create_linear_layer(self.input_dim, self.n_dim, bias=False),
+            nn.LeakyReLU(0.1),
+            self._create_linear_layer(self.n_dim, len(emb_list), bias=False),
+        ]
+        self.gate_emb = nn.Sequential(*gate_emb_layer_list)
+
+    def add_router_for_skipping(self) -> None:
+
+        self.output_dim = int(len(self.config.layer_ranking_list) + 1)
+
+        gate_skip_mlp_layer_list = [
+            self._create_linear_layer(self.input_dim, self.n_dim, bias=False),
+            nn.LeakyReLU(0.1),
+            self._create_linear_layer(self.n_dim, self.output_dim, bias=False),
+        ]
+
+        self.gate_skip_layer = nn.Sequential(*gate_skip_mlp_layer_list)
+
+    def mlp_forward(self, curr_iteration, budget_tensor, device, dtype, tau, hard_sample):
+
+        # Megatron's `[0]` after each Linear call unpacks TEColumnParallelLinear's
+        # (output, bias) tuple return. nn.Linear returns a tensor directly, so
+        # the unpacks are dropped here.
+        router_mlp_logits1 = self.gate_mlp[0](budget_tensor)
+        router_mlp_logits2 = self.gate_mlp[1](router_mlp_logits1)
+        router_mlp_logits = self.gate_mlp[2](router_mlp_logits2).flatten()
+        # Megatron leaves `scale` unbound when `self.scaler is None`, then
+        # references it unconditionally below → NameError for the common
+        # scaler-off configuration. Default to 1.0 so the `scale * ...`
+        # multiplications become no-ops in that case.
+        scale = 1.0
+        if self.scaler is not None:
+            scale = self.scaler[curr_iteration].to(device=device, dtype=dtype)
+        if self.config.flex_hetero_ffn:
+            mlp_n = len(self.config.mlp_int_list)
+            router_mlp_logits = router_mlp_logits.reshape(-1, mlp_n)
+            if self.config.normalize_router_logits:
+                router_mlp_logits = (
+                    scale
+                    * router_mlp_logits
+                    / router_mlp_logits.std(dim=1, keepdim=True).clamp(min=1e-6)
+                )
+            else:
+                router_mlp_logits = scale * router_mlp_logits
+            router_mlp_logits = self._dp_gumbel_softmax(
+                router_mlp_logits, tau=tau, hard=hard_sample, curr_iteration=curr_iteration
+            )
+            _, choices_mlp = torch.topk(router_mlp_logits, 1, dim=-1)
+            return (
+                router_mlp_logits,
+                [self.config.mlp_int_list[i] for i in choices_mlp.flatten().tolist()],
+            )
+        else:
+            if self.config.normalize_router_logits:
+                # Std-normalize only when there's actually >1 choice; with a
+                # single choice the std is 0 and the routing is trivial, so we
+                # skip both the scale and the normalization (consistent with
+                # the no-op semantics of a single-choice axis).
+                if len(self.config.mlp_int_list) > 1:
+                    router_mlp_logits = (
+                        scale
+                        * router_mlp_logits
+                        / router_mlp_logits.std(dim=0, keepdim=True).clamp(min=1e-6)
+                    )
+            else:
+                router_mlp_logits = scale * router_mlp_logits
+            router_mlp_logits = self._dp_gumbel_softmax(
+                router_mlp_logits, tau=tau, hard=hard_sample, curr_iteration=curr_iteration
+            )
+            _, choices_mlp = torch.topk(router_mlp_logits, 1, dim=-1)
+            return (router_mlp_logits, self.config.mlp_int_list[choices_mlp.item()])
+
+    def emb_forward(self, curr_iteration, budget_tensor, device, dtype, tau, hard_sample):
+
+        # `[0]` unpacks removed vs Megatron — see comment in mlp_forward.
+        router_emb_logits1 = self.gate_emb[0](budget_tensor)
+        router_emb_logits2 = self.gate_emb[1](router_emb_logits1)
+        router_emb_logits = self.gate_emb[2](router_emb_logits2).flatten()
+        if self.scaler is not None:
+            scale = self.scaler[curr_iteration].to(device=device, dtype=dtype)
+            router_emb_logits = scale * router_emb_logits
+
+        router_emb_logits = self._dp_gumbel_softmax(
+            router_emb_logits, tau=tau, hard=hard_sample, curr_iteration=curr_iteration
+        )
+        _, choices_emb = torch.topk(router_emb_logits, 1, dim=-1)
+
+        return (router_emb_logits, self.config.emb_int_list[choices_emb.item()])
+
+    def skipping_forward(self, curr_iteration, budget_tensor, device, dtype, tau, hard_sample):
+
+        # for layer skipping, skipping MLP layers
+        # `[0]` unpacks removed vs Megatron — see comment in mlp_forward.
+        router_skip_layer_logits1 = self.gate_skip_layer[0](budget_tensor)
+        router_skip_layer_logits2 = self.gate_skip_layer[1](router_skip_layer_logits1)
+        router_skip_layer_logits = self.gate_skip_layer[2](router_skip_layer_logits2).flatten()
+        router_skip_layer_logits = torch.repeat_interleave(
+            router_skip_layer_logits, repeats=1, dim=0
+        )
+        if self.scaler is not None:
+            router_skip_layer_logits = router_skip_layer_logits * self.scaler[
+                curr_iteration
+            ].to(device=device, dtype=dtype)
+
+        router_skip_layer_logits = self._dp_gumbel_softmax(
+            router_skip_layer_logits, tau=tau, hard=hard_sample, curr_iteration=curr_iteration
+        )
+        _, choices_skip_layer = torch.topk(router_skip_layer_logits, 1, dim=-1)
+        if choices_skip_layer.item() != 0:
+            selected_to_drop = self.config.layer_ranking_list[: choices_skip_layer.item()]
+            choices_skip_layer = torch.zeros(self.config.num_layers).to(device=device, dtype=dtype)
+            choices_skip_layer[selected_to_drop] = 1
+        else:
+            choices_skip_layer = torch.zeros(self.config.num_layers).to(device=device, dtype=dtype)
+        return (router_skip_layer_logits, choices_skip_layer)
+
+    def get_curr_tau(self, curr_iteration: int) -> torch.Tensor:
+        """Gumbel-softmax temperature: `tau_init * tau_decay ** curr_iteration`."""
+        return self.config.tau_init * torch.pow(
+            torch.tensor(self.config.tau_decay), curr_iteration
+        )
+
+    def add_scaler_schedule(self) -> None:
+        """Precompute an optional linear logit-scaler over training iterations.
+
+        Off unless both scaler endpoints and `train_iters` are set. When on,
+        each forward multiplies its axis logits by `scaler[iteration]` before
+        Gumbel-softmax — sharpens routing over training in a more controlled
+        way than tau decay alone.
+        """
+        cfg = self.config
+        if (
+            cfg.linear_scaler_start is None
+            or cfg.linear_scaler_end is None
+            or cfg.train_iters is None
+        ):
+            self.scaler = None
+            return
+        self.scaler = torch.linspace(
+            start=cfg.linear_scaler_start,
+            end=cfg.linear_scaler_end,
+            steps=cfg.train_iters,
+        )
+
+    def forward(self, budget, curr_iteration):
+
+        hard_sample = random.random() > self.hard_sample_th
+
+        tau = self.get_curr_tau(curr_iteration)
+
+        device, dtype = next(self.parameters()).device, next(self.parameters()).dtype
+
+        if budget in self.budget_map.keys():
+            budget_tensor = F.one_hot(
+                self.budget_map[budget], len(self.config.budget_list)
+            ).to(device=device, dtype=dtype)
+        elif budget == 1.0:
+            # Requested full model but 1.0 isn't a trained budget — fall back
+            # to the largest configured budget. Using max() instead of [0]
+            # makes this independent of budget_list ordering.
+            budget_tensor = F.one_hot(
+                self.budget_map[max(self.budget_map.keys())],
+                len(self.config.budget_list),
+            ).to(device=device, dtype=dtype)
+        else:
+            # budget_list is enforced descending by FlextronConfig.__attrs_post_init__.
+            # We re-sort ascending locally for bucketize, then flip(0) below to
+            # land back in the descending one-hot coordinate system the router
+            # was trained against.
+            budget_values = torch.tensor(
+                sorted(self.config.budget_list), device=device, dtype=dtype
+            )
+            budget_t = torch.as_tensor(budget, device=device, dtype=dtype)
+
+            # idx2 = first index where budget_values[idx] > budget (right=False gives >= behavior with floats)
+            idx2 = torch.bucketize(budget_t, budget_values, right=False)
+            # Clamp to valid interior so we always have a left neighbor
+            idx2 = idx2.clamp(min=1, max=len(self.config.budget_list) - 1)
+            idx1 = idx2 - 1
+
+            b1 = budget_values.index_select(0, idx1.to(torch.long))
+            b2 = budget_values.index_select(0, idx2.to(torch.long))
+            denom = b2 - b1  # .clamp_min(1e-12)
+            weight = (budget_t - b1) / denom  # in [0,1] when budget is between b1 and b2
+
+            num_classes = len(self.config.budget_list)
+            one_hot_1 = F.one_hot(
+                idx1.to(torch.long), num_classes=num_classes
+            ).to(device=device, dtype=dtype)
+            one_hot_2 = F.one_hot(
+                idx2.to(torch.long), num_classes=num_classes
+            ).to(device=device, dtype=dtype)
+
+            # If weight is scalar, broadcasting works; if vector, it blends per-sample
+            budget_tensor = (1 - weight).unsqueeze(-1) * one_hot_1 + weight.unsqueeze(
+                -1
+            ) * one_hot_2
+            budget_tensor = budget_tensor.squeeze(0).flip(0)
+
+        budget_tensor = budget_tensor.unsqueeze(0)
+        mlp_forward_outputs = self.mlp_forward(
+            curr_iteration, budget_tensor, device, dtype, tau, hard_sample
+        )
+
+        if self.config.add_skipping:
+            skipping_forward_outputs = self.skipping_forward(
+                curr_iteration, budget_tensor, device, dtype, tau, hard_sample
+            )
+        else:
+            skipping_forward_outputs = None
+
+        emb_forward_outputs = self.emb_forward(
+            curr_iteration, budget_tensor, device, dtype, tau, hard_sample
+        )
+        self.fwd_pass_count += 1
+        # 3-tuple return (Megatron returns 5-tuple including mamba + moe_expert;
+        # we drop those two positions rather than pad with None).
+        return (
+            mlp_forward_outputs,
+            skipping_forward_outputs,
+            emb_forward_outputs,
+        )
diff --git a/tests/elastification/__init__.py b/tests/elastification/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/elastification/test_router.py b/tests/elastification/test_router.py
new file mode 100644
index 0000000..f1e89b9
--- /dev/null
+++ b/tests/elastification/test_router.py
@@ -0,0 +1,405 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Unit tests for `fastgen.methods.elastification.router.FlextronRouter`.
+
+Covers:
+  - Construction: config wiring, gate MLP shapes + init, DP awareness, scaler
+  - Forward path: trained-budget lookup, `budget=1.0` fallback, interpolation,
+    fwd_pass_count increment
+  - DP-seeded Gumbel-softmax: RNG isolation, determinism, rank/iter/fwd-count
+    sensitivity
+  - Tau decay schedule
+  - Gradient flow (with non-invariant loss — softmax outputs sum to 1)
+  - add_skipping=True path (skip gate + per-layer mask)
+  - flex_hetero_ffn=True path (per-block choice list)
+  - Scaler-off default config (was a Megatron NameError we fixed)
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from fastgen.methods.elastification.config import FlextronConfig
+from fastgen.methods.elastification.router import FlextronRouter
+
+
+# ─── Fixtures ────────────────────────────────────────────────────────────
+
+
+@pytest.fixture
+def basic_config() -> FlextronConfig:
+    """Config with elasticity on both axes, scaler off (the default v1 recipe)."""
+    return FlextronConfig(
+        emb_int_list=[5120, 3840, 2560, 1280],
+        mlp_int_list=[13824, 10368, 6912, 3456],
+        budget_list=[1.0, 0.75, 0.5, 0.25],
+        router_std=0.1,
+    )
+
+
+@pytest.fixture
+def scaler_config() -> FlextronConfig:
+    """Config identical to `basic_config` but with the linear scaler schedule on."""
+    return FlextronConfig(
+        emb_int_list=[5120, 3840, 2560, 1280],
+        mlp_int_list=[13824, 10368, 6912, 3456],
+        budget_list=[1.0, 0.75, 0.5, 0.25],
+        router_std=0.1,
+        linear_scaler_start=1.0,
+        linear_scaler_end=10.0,
+        train_iters=1000,
+    )
+
+
+@pytest.fixture
+def skipping_config() -> FlextronConfig:
+    """Config with `add_skipping=True`."""
+    return FlextronConfig(
+        emb_int_list=[5120, 3840, 2560, 1280],
+        mlp_int_list=[13824, 10368, 6912, 3456],
+        budget_list=[1.0, 0.75, 0.5],
+        add_skipping=True,
+        layer_ranking_list=[0, 5, 10, 15, 20],
+        num_layers=30,
+    )
+
+
+@pytest.fixture
+def hetero_config() -> FlextronConfig:
+    """Config with `flex_hetero_ffn=True`."""
+    return FlextronConfig(
+        emb_int_list=[5120, 3840],
+        mlp_int_list=[13824, 10368, 6912],
+        budget_list=[1.0, 0.5],
+        flex_hetero_ffn=True,
+        num_mlp_blocks=30,
+    )
+
+
+# ─── Construction ────────────────────────────────────────────────────────
+
+
+class TestConstruction:
+    def test_default_construction_succeeds(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert r.config is basic_config
+
+    def test_input_dim_matches_budget_list_length(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert r.input_dim == len(basic_config.budget_list)
+
+    def test_n_dim_matches_router_inter_dim(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert r.n_dim == basic_config.router_inter_dim
+
+    def test_budget_map_populated(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert set(r.budget_map.keys()) == set(basic_config.budget_list)
+        # Descending order: index 0 = largest budget.
+        for i, b in enumerate(basic_config.budget_list):
+            assert r.budget_map[b].item() == i
+
+    def test_dp_awareness_defaults(self, basic_config):
+        r = FlextronRouter(basic_config)
+        # In an uninitialized-dist context, rank=0 and size=1.
+        assert r.dp_rank == 0
+        assert r.dp_size == 1
+
+    def test_fwd_pass_count_starts_at_zero(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert r.fwd_pass_count == 0
+
+    def test_hard_sample_th_copied_from_config(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert r.hard_sample_th == basic_config.hard_sample_th
+
+    def test_gate_mlp_shape(self, basic_config):
+        r = FlextronRouter(basic_config)
+        # 3-element Sequential: Linear(input_dim, n_dim) → LeakyReLU → Linear(n_dim, len(mlp_int_list))
+        assert r.gate_mlp[0].weight.shape == (r.n_dim, r.input_dim)
+        assert r.gate_mlp[2].weight.shape == (len(basic_config.mlp_int_list), r.n_dim)
+
+    def test_gate_emb_shape(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert r.gate_emb[0].weight.shape == (r.n_dim, r.input_dim)
+        assert r.gate_emb[2].weight.shape == (len(basic_config.emb_int_list), r.n_dim)
+
+    def test_gates_have_no_bias(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert r.gate_mlp[0].bias is None
+        assert r.gate_mlp[2].bias is None
+        assert r.gate_emb[0].bias is None
+        assert r.gate_emb[2].bias is None
+
+    def test_weight_init_std_close_to_router_std(self, basic_config):
+        r = FlextronRouter(basic_config)
+        # 500-dim weight, N(0, 0.1) → sample std within ~7% of target 99% of the time
+        assert abs(r.gate_emb[0].weight.std().item() - basic_config.router_std) < 0.02
+
+    def test_scaler_none_by_default(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert r.scaler is None
+
+    def test_scaler_built_when_configured(self, scaler_config):
+        r = FlextronRouter(scaler_config)
+        assert r.scaler is not None
+        assert r.scaler.shape == (scaler_config.train_iters,)
+        assert r.scaler[0].item() == pytest.approx(scaler_config.linear_scaler_start)
+        assert r.scaler[-1].item() == pytest.approx(scaler_config.linear_scaler_end)
+
+
+# ─── Forward at various budgets ──────────────────────────────────────────
+
+
+class TestForward:
+    def test_forward_at_trained_budget_returns_three_tuple(self, basic_config):
+        r = FlextronRouter(basic_config)
+        out = r.forward(budget=0.75, curr_iteration=0)
+        assert len(out) == 3
+
+    def test_mlp_out_shape_matches_choice_list(self, basic_config):
+        r = FlextronRouter(basic_config)
+        mlp_out, _, _ = r.forward(budget=0.75, curr_iteration=0)
+        assert mlp_out[0].shape == (len(basic_config.mlp_int_list),)
+
+    def test_emb_out_shape_matches_choice_list(self, basic_config):
+        r = FlextronRouter(basic_config)
+        _, _, emb_out = r.forward(budget=0.75, curr_iteration=0)
+        assert emb_out[0].shape == (len(basic_config.emb_int_list),)
+
+    def test_hard_choice_is_valid_member_of_int_list(self, basic_config):
+        r = FlextronRouter(basic_config)
+        mlp_out, _, emb_out = r.forward(budget=0.75, curr_iteration=0)
+        assert mlp_out[1] in basic_config.mlp_int_list
+        assert emb_out[1] in basic_config.emb_int_list
+
+    def test_skip_output_none_when_add_skipping_false(self, basic_config):
+        r = FlextronRouter(basic_config)
+        _, skip_out, _ = r.forward(budget=0.75, curr_iteration=0)
+        assert skip_out is None
+
+    @pytest.mark.parametrize("budget", [1.0, 0.75, 0.5, 0.25])
+    def test_forward_at_each_trained_budget(self, basic_config, budget):
+        r = FlextronRouter(basic_config)
+        out = r.forward(budget=budget, curr_iteration=0)
+        assert out[0][0].shape == (len(basic_config.mlp_int_list),)
+        assert out[2][0].shape == (len(basic_config.emb_int_list),)
+
+    def test_forward_at_interpolated_budget(self, basic_config):
+        r = FlextronRouter(basic_config)
+        # 0.625 is not in [1.0, 0.75, 0.5, 0.25] → interpolation branch
+        out = r.forward(budget=0.625, curr_iteration=100)
+        assert out[0][1] in basic_config.mlp_int_list
+
+    def test_forward_at_budget_1_when_1_not_in_list(self):
+        """Fallback branch: `budget=1.0` requested but not in `budget_list`."""
+        cfg = FlextronConfig(
+            emb_int_list=[1024, 512],
+            mlp_int_list=[4096, 2048],
+            budget_list=[0.75, 0.5, 0.25],
+        )
+        r = FlextronRouter(cfg)
+        # Should fall back to largest configured budget (0.75) without crashing.
+        out = r.forward(budget=1.0, curr_iteration=0)
+        assert out[0][1] in cfg.mlp_int_list
+
+    def test_fwd_pass_count_increments_each_forward(self, basic_config):
+        r = FlextronRouter(basic_config)
+        n0 = r.fwd_pass_count
+        _ = r.forward(budget=0.75, curr_iteration=0)
+        _ = r.forward(budget=0.5, curr_iteration=1)
+        _ = r.forward(budget=0.75, curr_iteration=2)
+        assert r.fwd_pass_count == n0 + 3
+
+
+# ─── DP-seeded Gumbel-softmax ────────────────────────────────────────────
+
+
+class TestDPGumbelSoftmax:
+    def test_cpu_rng_state_restored(self, basic_config):
+        r = FlextronRouter(basic_config)
+        torch.manual_seed(1234)
+        x = torch.randn(4)  # advance RNG so `before` isn't the post-seed state
+        before = torch.get_rng_state()
+        _ = r._dp_gumbel_softmax(x, tau=0.5, hard=False, curr_iteration=100)
+        after = torch.get_rng_state()
+        assert torch.equal(before, after), "CPU RNG state must be restored"
+
+    def test_same_seed_produces_identical_sample(self, basic_config):
+        r = FlextronRouter(basic_config)
+        logits = torch.randn(4)
+        r.dp_rank = 0
+        r.fwd_pass_count = 0
+        s1 = r._dp_gumbel_softmax(logits.clone(), tau=1.0, hard=False, curr_iteration=42)
+        r.dp_rank = 0
+        r.fwd_pass_count = 0
+        s2 = r._dp_gumbel_softmax(logits.clone(), tau=1.0, hard=False, curr_iteration=42)
+        assert torch.equal(s1, s2)
+
+    def test_different_dp_rank_produces_different_sample(self, basic_config):
+        r = FlextronRouter(basic_config)
+        logits = torch.randn(4)
+        r.dp_rank = 0
+        r.fwd_pass_count = 0
+        s0 = r._dp_gumbel_softmax(logits.clone(), tau=1.0, hard=False, curr_iteration=42)
+        r.dp_rank = 1
+        r.fwd_pass_count = 0
+        s1 = r._dp_gumbel_softmax(logits.clone(), tau=1.0, hard=False, curr_iteration=42)
+        assert not torch.equal(s0, s1)
+
+    def test_different_iteration_produces_different_sample(self, basic_config):
+        r = FlextronRouter(basic_config)
+        logits = torch.randn(4)
+        r.dp_rank = 0
+        r.fwd_pass_count = 0
+        s_iter42 = r._dp_gumbel_softmax(logits.clone(), tau=1.0, hard=False, curr_iteration=42)
+        s_iter43 = r._dp_gumbel_softmax(logits.clone(), tau=1.0, hard=False, curr_iteration=43)
+        assert not torch.equal(s_iter42, s_iter43)
+
+    def test_different_fwd_pass_count_produces_different_sample(self, basic_config):
+        r = FlextronRouter(basic_config)
+        logits = torch.randn(4)
+        r.dp_rank = 0
+        r.fwd_pass_count = 0
+        s_a = r._dp_gumbel_softmax(logits.clone(), tau=1.0, hard=False, curr_iteration=42)
+        r.fwd_pass_count = 1
+        s_b = r._dp_gumbel_softmax(logits.clone(), tau=1.0, hard=False, curr_iteration=42)
+        assert not torch.equal(s_a, s_b)
+
+
+# ─── Tau schedule ────────────────────────────────────────────────────────
+
+
+class TestTauDecay:
+    def test_tau_at_iteration_zero_equals_tau_init(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert r.get_curr_tau(0).item() == pytest.approx(basic_config.tau_init)
+
+    def test_tau_decreases_monotonically(self, basic_config):
+        r = FlextronRouter(basic_config)
+        taus = [r.get_curr_tau(i).item() for i in (0, 1000, 10000, 30000)]
+        assert taus == sorted(taus, reverse=True)
+
+    def test_tau_approaches_zero_at_large_iteration(self, basic_config):
+        r = FlextronRouter(basic_config)
+        # 0.9999 ** 30000 ≈ 0.0498
+        assert r.get_curr_tau(30000).item() < 0.06
+
+
+# ─── Gradient flow ───────────────────────────────────────────────────────
+
+
+class TestGradientFlow:
+    def test_gradient_flows_to_all_router_params(self, basic_config):
+        """
+        Uses a non-invariant loss.  A common pitfall: `sum(softmax(x)) = 1`
+        identically, so `mlp_probs.sum() + emb_probs.sum()` is a constant
+        w.r.t. the logits and produces zero gradient.  We use an
+        arange-weighted sum instead.
+        """
+        r = FlextronRouter(basic_config)
+        for p in r.parameters():
+            p.grad = None
+
+        mlp_out, _, emb_out = r.forward(budget=0.75, curr_iteration=0)
+        w_mlp = torch.arange(len(basic_config.mlp_int_list), dtype=torch.float32)
+        w_emb = torch.arange(len(basic_config.emb_int_list), dtype=torch.float32)
+        loss = (mlp_out[0] * w_mlp).sum() + (emb_out[0] * w_emb).sum()
+        loss.backward()
+
+        for name, p in r.named_parameters():
+            assert p.grad is not None, f"{name} missing gradient"
+            assert p.grad.abs().sum() > 0, f"{name} received zero gradient"
+
+
+# ─── add_skipping branch ─────────────────────────────────────────────────
+
+
+class TestAddSkipping:
+    def test_gate_skip_layer_built_when_add_skipping_true(self, skipping_config):
+        r = FlextronRouter(skipping_config)
+        assert hasattr(r, "gate_skip_layer")
+
+    def test_gate_skip_layer_not_built_when_add_skipping_false(self, basic_config):
+        r = FlextronRouter(basic_config)
+        assert not hasattr(r, "gate_skip_layer")
+
+    def test_skip_gate_output_dim(self, skipping_config):
+        r = FlextronRouter(skipping_config)
+        # Last-layer output dim = len(layer_ranking_list) + 1 (one extra "skip 0 layers" option)
+        assert r.gate_skip_layer[2].weight.shape[0] == len(skipping_config.layer_ranking_list) + 1
+
+    def test_skip_output_non_none_when_add_skipping_true(self, skipping_config):
+        r = FlextronRouter(skipping_config)
+        _, skip_out, _ = r.forward(budget=0.75, curr_iteration=0)
+        assert skip_out is not None
+
+    def test_skip_probs_shape(self, skipping_config):
+        r = FlextronRouter(skipping_config)
+        _, skip_out, _ = r.forward(budget=0.75, curr_iteration=0)
+        soft_probs, _ = skip_out
+        assert soft_probs.shape == (len(skipping_config.layer_ranking_list) + 1,)
+
+    def test_skip_mask_shape_matches_num_layers(self, skipping_config):
+        r = FlextronRouter(skipping_config)
+        _, skip_out, _ = r.forward(budget=0.75, curr_iteration=0)
+        _, skip_mask = skip_out
+        assert skip_mask.shape == (skipping_config.num_layers,)
+
+    def test_skip_mask_values_are_zero_or_one(self, skipping_config):
+        r = FlextronRouter(skipping_config)
+        _, skip_out, _ = r.forward(budget=0.5, curr_iteration=0)
+        _, skip_mask = skip_out
+        # Per Megatron's construction: exactly 0 or 1 per layer.
+        assert set(skip_mask.unique().tolist()).issubset({0.0, 1.0})
+
+
+# ─── flex_hetero_ffn branch ──────────────────────────────────────────────
+
+
+class TestFlexHeteroFfn:
+    def test_gate_mlp_output_dim_scales_with_num_blocks(self, hetero_config):
+        r = FlextronRouter(hetero_config)
+        expected_output_dim = len(hetero_config.mlp_int_list) * hetero_config.num_mlp_blocks
+        assert r.gate_mlp[2].weight.shape[0] == expected_output_dim
+
+    def test_hetero_mlp_choice_is_a_list(self, hetero_config):
+        r = FlextronRouter(hetero_config)
+        mlp_out, _, _ = r.forward(budget=0.5, curr_iteration=0)
+        assert isinstance(mlp_out[1], list)
+
+    def test_hetero_choice_list_has_num_mlp_blocks_entries(self, hetero_config):
+        r = FlextronRouter(hetero_config)
+        mlp_out, _, _ = r.forward(budget=0.5, curr_iteration=0)
+        assert len(mlp_out[1]) == hetero_config.num_mlp_blocks
+
+    def test_hetero_choices_are_all_valid(self, hetero_config):
+        r = FlextronRouter(hetero_config)
+        mlp_out, _, _ = r.forward(budget=0.5, curr_iteration=0)
+        for c in mlp_out[1]:
+            assert c in hetero_config.mlp_int_list
+
+
+# ─── Scaler-off default config (Megatron NameError we fixed) ────────────
+
+
+class TestScalerOffPath:
+    """The default config (`scaler=None`, `normalize_router_logits=False`) used
+    to raise `NameError: local variable 'scale' referenced before assignment`
+    in Megatron's `mlp_forward`.  We default `scale = 1.0` so the multiplication
+    is a no-op.  See PLAN.md "Known Megatron quirks"."""
+
+    def test_default_config_forward_does_not_raise(self, basic_config):
+        r = FlextronRouter(basic_config)
+        r.forward(budget=0.75, curr_iteration=0)  # must not raise NameError
+
+    def test_hetero_forward_with_scaler_off_does_not_raise(self, hetero_config):
+        r = FlextronRouter(hetero_config)
+        r.forward(budget=0.5, curr_iteration=0)  # must not raise NameError
+
+    def test_skipping_forward_with_scaler_off_does_not_raise(self, skipping_config):
+        r = FlextronRouter(skipping_config)
+        r.forward(budget=0.75, curr_iteration=0)  # must not raise NameError