freesolo-flash-dev 0.2.42__tar.gz → 0.2.43__tar.gz

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.43}/Dockerfile.worker

@@ -69,7 +69,12 @@ RUN pip install --no-cache-dir \
 # miscomputed with Triton>=3.4 (fla #640) — the fix is fla's TILELANG backend (correct there), so
 # we bake the pinned tilelang==0.1.11 + apache-tvm-ffi==0.1.11 (0.1.12 double-registers TVM-FFI ->
 # `import tilelang` aborts) and KEEP fla on every arch; flash.engine.worker._ensure_fla_fastpath_on_hopper
-# ensures the tilelang backend is live on sm90 before any model import. Validated win vs the old
+# ensures the tilelang backend is live on sm90 before any model import. The INVERSE holds on
+# B200/sm100: THERE the tilelang chunk_bwd_dqkwg miscomputes gradients while the Triton path is
+# correct — and the pinned fla default-enables tilelang wherever it imports (i.e. everywhere,
+# since it is baked here for Hopper), so the worker opts sm100 out via FLA_TILELANG=0
+# (engine.worker.perf._force_fla_triton_gdn_on_sm100; upstream default-gates tilelang to
+# Hopper-only since fla #975). Validated win vs the old
 # pure-PyTorch fallback (H100 SXM, Qwen3.5 hidden-2560 LoRA): seq4096 4.2x faster / 1.6x less mem,
 # seq16384 12.6x / 1.9x, loss matches to 1.8e-4. fla is from git: the PyPI wheel is a broken stub
 # (no fla.modules). The git URL is PINNED to a commit (not the moving default branch) for

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.43}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.42
+Version: 0.2.43
 Summary: Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI
 Project-URL: Homepage, https://github.com/freesolo-co/flash
 Project-URL: Repository, https://github.com/freesolo-co/flash

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.43}/flash/engine/worker/__init__.py

@@ -100,6 +100,7 @@ from flash.engine.worker.perf import (
     _estimate_params,
     _flash_attn_3_available,
     _flash_attn_available,
+    _force_fla_triton_gdn_on_sm100,
     _GpuPeakSampler,
     _liger_default_for_model,
     _memory_mode,
@@ -108,6 +109,7 @@ from flash.engine.worker.perf import (
     _peak_gpu_gb,
     _remove_fla_from_disk,
     _reset_peak_gpu,
+    _restrict_fla_gdn_autotune_on_blackwell,
     _sdpa_cudnn_ctx,
     free_gpu,
     fused_optim_name,
@@ -265,9 +267,15 @@ def main():
                 raise RetriableInfraError(
                     f"DONE present but metrics.json unreadable after retries (transient HF): {last_err}"
                 )
+        # BEFORE any model import / fla dispatch: on sm100 the baked tilelang GDN backend
+        # computes wrong gradients — opt out so fla uses its (correct-there) Triton path.
+        _force_fla_triton_gdn_on_sm100()
         _ensure_fla_fastpath_on_hopper()
         # Must run AFTER fla fast path (may reinstall tilelang) and BEFORE model/vLLM import.
         _neutralize_tilelang_cudart_stub()
+        # AFTER the fla fast path (which may (re)install fla), BEFORE any model import / GDN
+        # launch: restrict fla's Blackwell GDN bwd autotune to grad-correct configs (fla #913).
+        _restrict_fla_gdn_autotune_on_blackwell()
         heartbeat("boot", gpu=gpu_diagnostics(include_torch=False))
         finalize_alloc_conf_for_sleep()
         load_mega_cache()
@@ -344,6 +352,7 @@ __all__ = [
     "_finalize",
     "_flash_attn_3_available",
     "_flash_attn_available",
+    "_force_fla_triton_gdn_on_sm100",
     "_grpo_is_no_op_failure",
     "_grpo_resume_already_complete",
     "_hf_upload",
@@ -359,6 +368,7 @@ __all__ = [
     "_remove_fla_from_disk",
     "_reset_peak_gpu",
     "_resolve_adapter_ref",
+    "_restrict_fla_gdn_autotune_on_blackwell",
     "_sdpa_cudnn_ctx",
     "assert_adapter_delta_nonzero",
     "assert_adapter_load_clean",

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.43}/flash/engine/worker/kernel_warmup.py

@@ -77,7 +77,9 @@ def _require_gpu():
         if not torch.cuda.is_available():
             _log("no CUDA device visible — kernel warmup must run on a GPU builder; nothing baked")
             return None
-        _log(f"GPU: {torch.cuda.get_device_name(0)} ({_torch_sm(torch)}), torch {torch.__version__}")
+        _log(
+            f"GPU: {torch.cuda.get_device_name(0)} ({_torch_sm(torch)}), torch {torch.__version__}"
+        )
         return torch
     except Exception as e:
         _log(f"torch unavailable ({e}); cannot warm kernels")
@@ -128,6 +130,13 @@ def warm_flash_attn(torch) -> bool:
 def warm_fla_gdn(torch) -> bool:
     """Compile flash-linear-attention's Gated-DeltaNet chunk kernels (Qwen3.5/3.6 hybrid path)."""
     try:
+        # sm100: tilelang's GDN bwd miscomputes grads — warm the Triton path the worker will run.
+        try:
+            from flash.engine.worker.perf import _force_fla_triton_gdn_on_sm100
+
+            _force_fla_triton_gdn_on_sm100()
+        except Exception as e:
+            _log(f"sm100 fla tilelang opt-out skipped: {e}")
         # fla #640: GDN chunk_bwd needs tilelang on Hopper; Triton path miscomputes/raises there.
         try:
             from flash.engine.worker.perf import _ensure_fla_fastpath_on_hopper
@@ -135,6 +144,14 @@ def warm_fla_gdn(torch) -> bool:
             _ensure_fla_fastpath_on_hopper()
         except Exception as e:
             _log(f"hopper fla fast-path setup skipped: {e}")
+        # fla #913: on Blackwell, warm the SAME restricted autotune space the worker boot will
+        # run (grad-correct configs only) so the baked cache holds the config that actually runs.
+        try:
+            from flash.engine.worker.perf import _restrict_fla_gdn_autotune_on_blackwell
+
+            _restrict_fla_gdn_autotune_on_blackwell()
+        except Exception as e:
+            _log(f"blackwell fla autotune restriction skipped: {e}")
         from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 
         b, h, t, d = 1, 4, 256, 64

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.43}/flash/engine/worker/perf/__init__.py

@@ -131,7 +131,7 @@ def _find_real_libcudart() -> str | None:
             with open("/proc/self/maps") as f:
                 for line in f:
                     if base in line and "/" in line:
-                        p = line[line.index("/"):].rstrip()
+                        p = line[line.index("/") :].rstrip()
                         if os.path.basename(p).startswith(base) and os.path.exists(p):
                             return os.path.realpath(p)
         except OSError:
@@ -208,6 +208,138 @@ def _neutralize_tilelang_cudart_stub() -> None:
         print(f"[worker] libcudart stub neutralize failed: {e}", flush=True)
 
 
+def _force_fla_triton_gdn_on_sm100() -> None:
+    """Opt fla OUT of its tilelang GDN backend on B200 (sm100) via ``FLA_TILELANG=0``.
+
+    The worker image bakes tilelang for the HOPPER fast path (fla #640), but the pinned fla
+    default-enables its tilelang backend WHEREVER tilelang imports. On sm100, tilelang's
+    ``chunk_bwd_dqkwg`` computes WRONG GRADIENTS — measured on a real B200 with this exact image
+    at the production Qwen3.5/3.6 GDN call shapes (H==HV after transformers' head repeat):
+    dq/dk rel-err ~0.72, dg ~1.28 vs the fp32 reference, deterministic, in bf16 AND fp32, while
+    the same shapes on H200 are correct to ~4e-3. This is the root cause of the B200 35B-A3B SFT
+    incident (grad_norm ~1e8 from the first logged step; loss flat at LR 2e-5/5e-5, collapse at
+    1e-4 — garbage gradient DIRECTIONS no LR can fix, while H200 trained healthily at 1e-4).
+
+    ``FLA_TILELANG=0`` flips fla's dispatch to its Triton chunk_bwd_dqkwg — verified correct on
+    sm100 at the same shapes (~4e-3). Upstream has since stopped default-enabling tilelang
+    off-Hopper (fla #975 gates the default to sm90 + Triton>=3.4) and an explicit env opt-out
+    wins on every fla version, so this env set reproduces upstream's current gate on our pin and
+    stays a correct no-op across a future pin bump.
+
+    * ``setdefault``: an explicitly pre-set FLA_TILELANG (e.g. testing a fixed tilelang) wins.
+    * sm100 only: sm90 NEEDS tilelang (fla #640, the fast-path installer owns it); sm89/sm120
+      train healthily today under the pin's default and upstream's next-pin gate flips them to
+      Triton anyway — no evidence to justify changing them from here.
+    * Env is read live per-dispatch by fla, but set it at boot BEFORE any model import so no
+      kernel ever launches under the old default.
+    """
+    try:
+        import torch
+
+        if not (torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0)):
+            return
+    except Exception:
+        return
+    prior = os.environ.get("FLA_TILELANG")
+    if prior is not None:
+        print(
+            f"[blackwell] FLA_TILELANG={prior!r} pre-set; respecting it "
+            "(default here is 0 on sm100: tilelang chunk_bwd_dqkwg miscomputes grads)",
+            flush=True,
+        )
+        return
+    os.environ["FLA_TILELANG"] = "0"
+    print(
+        "[blackwell] sm100 (B200): FLA_TILELANG=0 -> fla GDN uses the Triton chunk_bwd_dqkwg "
+        "(tilelang backend computes wrong gradients on sm100; measured dq/dk ~0.72, dg ~1.28 "
+        "rel-err; upstream default-gates tilelang to Hopper since fla #975)",
+        flush=True,
+    )
+
+
+def _restrict_fla_gdn_autotune_on_blackwell() -> None:
+    """Restrict fla's GDN ``prepare_wy_repr_bwd_kernel`` autotune space on Blackwell (sm100/sm120).
+
+    Defense-in-depth for the Triton GDN backward path — the path sm100 defaults to once
+    ``_force_fla_triton_gdn_on_sm100`` opts out of tilelang. Upstream reports that Blackwell can
+    select unstable Triton configs for this kernel during autotuning (fla #913) and restricted
+    the space to the B200-validated config (num_warps=2, num_stages=4) in fla #1000 (ac6c648) —
+    NEWER than the worker's pinned fla SHA — so apply the same restriction in-process at boot,
+    before any GDN kernel launch. (Not the 35B-A3B incident root cause — that was tilelang, see
+    above — but there is no reason to keep autotuning over configs upstream deemed unstable on
+    the exact card we train on.)
+
+    * No-op on non-Blackwell arches, when fla is absent (pure-PyTorch delta path), or when the
+      pinned fla already carries the upstream fix (filtering an already-restricted space).
+    * Fail-CLOSED like the Hopper fast path: if fla is present on Blackwell but the restriction
+      cannot be applied (autotuner layout changed — e.g. an unreviewed pin bump), physically
+      remove fla so transformers falls back to the pure-PyTorch delta rule. Slow-but-correct
+      beats a run that burns GPU-hours computing garbage gradients.
+    * sm120 is included to match upstream's IS_NVIDIA_BLACKWELL gate: 5090 smokes trained fine on
+      the unrestricted space, but that is autotune luck (the winning config varies per shape/run),
+      not evidence of safety.
+    """
+    try:
+        import torch
+
+        if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] in (10, 12)):
+            return
+    except Exception:
+        return
+    try:
+        import importlib.util
+
+        if importlib.util.find_spec("fla") is None:
+            return  # no fla -> transformers already uses the pure-PyTorch delta path
+    except Exception:
+        return
+
+    def _fail_closed(reason: str) -> None:
+        _removed, _still = _remove_fla_from_disk()
+        print(
+            f"[blackwell] fla GDN autotune restriction FAILED ({reason}) -> DISABLING fla "
+            f"(removed {len(_removed)} copy(ies); still_importable={_still}); "
+            "pure-PyTorch delta fallback (fla #913: unrestricted autotune miscomputes grads)",
+            flush=True,
+        )
+
+    try:
+        from fla.ops.gated_delta_rule import wy_fast
+
+        tuner = getattr(wy_fast, "prepare_wy_repr_bwd_kernel", None)
+        for _ in range(8):  # unwrap decorator layers (heuristics/cache) down to the Autotuner
+            if tuner is None or hasattr(tuner, "configs"):
+                break
+            tuner = getattr(tuner, "fn", None)
+        configs = getattr(tuner, "configs", None)
+        if not configs:
+            _fail_closed("prepare_wy_repr_bwd_kernel autotuner not found")
+            return
+        keep = [
+            c
+            for c in configs
+            if getattr(c, "num_warps", None) == 2 and getattr(c, "num_stages", None) == 4
+        ]
+        if not keep:
+            _fail_closed("no B200-validated config (num_warps=2/num_stages=4) in autotune space")
+            return
+        if len(keep) == len(configs):
+            print(
+                "[blackwell] fla GDN autotune already restricted (pinned fla carries fla #1000)",
+                flush=True,
+            )
+            return
+        tuner.configs = keep
+        print(
+            f"[blackwell] fla GDN prepare_wy_repr_bwd autotune restricted {len(configs)} -> "
+            f"{len(keep)} config(s) (num_warps=2/num_stages=4; fla #913/#1000 grad miscompute)",
+            flush=True,
+        )
+    except Exception as e:
+        with contextlib.suppress(Exception):
+            _fail_closed(f"{type(e).__name__}: {e}")
+
+
 def _ensure_fla_fastpath_on_hopper() -> None:
     """Enable fla's tilelang GDN backend on Hopper (sm90); no-op elsewhere.
 
@@ -334,6 +466,7 @@ __all__ = [
     "_flash_attn_3_available",
     "_flash_attn_available",
     "_float_or_none",
+    "_force_fla_triton_gdn_on_sm100",
     "_int_or_none",
     "_liger_default_for_model",
     "_memory_mode",
@@ -344,6 +477,7 @@ __all__ = [
     "_query_nvidia_processes",
     "_remove_fla_from_disk",
     "_reset_peak_gpu",
+    "_restrict_fla_gdn_autotune_on_blackwell",
     "_round_gb_from_mib",
     "_sdpa_cudnn_ctx",
     "free_gpu",

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.43}/flash/engine/worker/sft.py

@@ -99,6 +99,21 @@ def _model_arch_dims(model_id: str) -> tuple[int, int]:
         return c_hidden, c_layers
 
 
+def select_sft_examples(train, max_examples, seed):
+    """Pick the SFT sample: the first ``max_examples`` rows of the dataset (file order), shuffled.
+
+    The slice happens BEFORE the shuffle so ``max_examples`` is a deterministic prefix fence,
+    not a random subsample. A train.jsonl that carries fully-labeled SFT rows first and
+    prompt-only (empty-output) GRPO rows after can cap SFT to the labeled head — an empty
+    completion can never be shuffled into the SFT sample and teach the model to emit nothing.
+    """
+    if max_examples > 0:
+        train = train[:max_examples]
+    rng = random.Random(seed)
+    rng.shuffle(train)
+    return train
+
+
 def run_sft():
     from datasets import Dataset
     from transformers import AutoTokenizer
@@ -122,12 +137,7 @@ def run_sft():
         val = getattr(_t, name, None) if _t else None
         return val if val is not None else default
 
-    train = env.dataset()
-    rng = random.Random(_w.SEED)
-    rng.shuffle(train)
-    max_examples = int(_train_opt("max_examples", 0) or 0)
-    if max_examples > 0:
-        train = train[:max_examples]
+    train = select_sft_examples(env.dataset(), int(_train_opt("max_examples", 0) or 0), _w.SEED)
     texts = []
     multiturn_targets = 0
     for ex in train:

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.43}/flash/providers/_worker.py

@@ -41,7 +41,9 @@ WORKER_DEPS = [
     "wandb>=0.17",
     # fla from git: PyPI wheel is a broken stub missing fla.modules. SHA-pinned for reproducibility;
     # keep in lockstep with Dockerfile.worker. fla kept on ALL arches — worker ensures tilelang
-    # backend on sm90 before model import (fla #640: chunk_bwd miscompute with Triton>=3.4 on Hopper).
+    # backend on sm90 before model import (fla #640: chunk_bwd miscompute with Triton>=3.4 on Hopper)
+    # and OPTS OUT of tilelang on sm100 (B200) where tilelang's chunk_bwd_dqkwg miscomputes grads
+    # (worker _force_fla_triton_gdn_on_sm100; upstream default-gates tilelang to Hopper since fla #975).
     "flash-linear-attention @ git+https://github.com/fla-org/flash-linear-attention.git@f0e213dbd8b5fb90c3c7eca869ac1706d5377139",
     # tilelang version-pinned in lockstep with Dockerfile.worker + perf.py runtime reinstall.
     "tilelang==0.1.11",

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.43}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "0.2.42"
+version = "0.2.43"
 description = "Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI"
 readme = "README.md"
 # RunPod Flash supports Python 3.11-3.12 (not 3.13 yet).
@@ -106,7 +106,7 @@ artifacts = ["flash/py.typed"]
 # auto-published by .github/workflows/publish-dev.yml: a push to `dev` whose version isn't yet on
 # PyPI publishes it; pushes that don't change it find it already published and no-op. So cutting a
 # dev-channel release = bump this (and [project].version) and merge to `dev`.
-version = "0.2.42"
+version = "0.2.43"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

freesolo_flash_dev-0.2.43/tests/test_sft_example_selection.py

@@ -0,0 +1,36 @@
+"""``max_examples`` must fence SFT to the HEAD of train.jsonl, slicing BEFORE the shuffle.
+
+A dataset may carry fully-labeled SFT rows first and prompt-only (empty-output) GRPO rows
+after; with shuffle-then-slice an empty completion gets sampled into SFT and trains the
+model to emit nothing. Regression for the shuffle-before-slice ordering in run_sft.
+"""
+
+from __future__ import annotations
+
+from flash.engine.worker.sft import select_sft_examples
+
+
+def test_max_examples_is_a_prefix_fence_not_a_random_subsample():
+    # rows 0..9 are labeled SFT rows; rows 10..19 are prompt-only GRPO rows (output="")
+    labeled = [{"input": f"q{i}", "output": f"a{i}"} for i in range(10)]
+    prompt_only = [{"input": f"p{i}", "output": ""} for i in range(10)]
+    picked = select_sft_examples(labeled + prompt_only, 10, seed=0)
+    assert len(picked) == 10
+    assert sorted(r["input"] for r in picked) == sorted(r["input"] for r in labeled)
+    # the property that actually matters: no empty completion leaks into the SFT sample
+    assert all(r["output"] for r in picked)
+
+
+def test_selection_is_seed_deterministic_but_still_shuffled():
+    rows = list(range(100))
+    a = select_sft_examples(list(rows), 50, seed=7)
+    b = select_sft_examples(list(rows), 50, seed=7)
+    assert a == b  # same seed -> same training order
+    assert sorted(a) == list(range(50))  # exactly the head of the file...
+    assert a != list(range(50))  # ...but shuffled for training
+
+
+def test_zero_or_unset_max_examples_keeps_the_full_dataset():
+    rows = list(range(20))
+    out = select_sft_examples(list(rows), 0, seed=3)
+    assert sorted(out) == rows

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.43}/tests/test_worker_stack.py

@@ -400,7 +400,9 @@ def test_heartbeat_rolls_back_slot_when_upload_reports_failure(monkeypatch):
     import flash.engine.worker as w
 
     monkeypatch.setattr(w, "_HB_MIN_INTERVAL_S", 0.0)
-    sentinel_last_upload = 123.0  # a prior successful-commit timestamp the failed retry must restore
+    sentinel_last_upload = (
+        123.0  # a prior successful-commit timestamp the failed retry must restore
+    )
     monkeypatch.setattr(w, "_HB_LAST_UPLOAD", sentinel_last_upload)
 
     calls = []
@@ -471,8 +473,12 @@ def test_critical_stages_wait_longer_for_upload_lock(monkeypatch):
     finally:
         hbmod._HB_UPLOAD_LOCK.release()
 
-    assert progress_wait < 0.3, f"progress stage should skip after the short timeout, waited {progress_wait:.2f}s"
-    assert critical_wait >= 0.3, f"critical stage should wait the long timeout, waited {critical_wait:.2f}s"
+    assert progress_wait < 0.3, (
+        f"progress stage should skip after the short timeout, waited {progress_wait:.2f}s"
+    )
+    assert critical_wait >= 0.3, (
+        f"critical stage should wait the long timeout, waited {critical_wait:.2f}s"
+    )
     assert critical_wait > progress_wait + 0.15
 
 
@@ -1019,7 +1025,9 @@ def test_tilelang_pin_is_consistent_and_pinned():
     perf_src = (root / "flash" / "engine" / "worker" / "perf" / "__init__.py").read_text()
     # perf/__init__.py builds the spec via an f-string `f"tilelang=={TILELANG_PIN}"`, so assert the constant.
     pm = re.search(r'TILELANG_PIN\s*=\s*"([0-9][0-9A-Za-z.\-]*)"', perf_src)
-    assert pm, "perf/__init__.py must define a pinned TILELANG_PIN constant for the runtime reinstall"
+    assert pm, (
+        "perf/__init__.py must define a pinned TILELANG_PIN constant for the runtime reinstall"
+    )
     assert pm.group(1) == pin, (
         f"perf/__init__.py TILELANG_PIN must match WORKER_DEPS (deps={pin}, perf={pm.group(1)})"
     )
@@ -1188,7 +1196,9 @@ def test_neutralize_never_loads_the_stub_into_the_process(tmp_path, monkeypatch)
     if not os.path.exists("/proc/self/maps"):
         import pytest
 
-        pytest.skip("/proc/self/maps unavailable (non-Linux); the loaded-mapping assertion needs it")
+        pytest.skip(
+            "/proc/self/maps unavailable (non-Linux); the loaded-mapping assertion needs it"
+        )
 
     pkg = tmp_path / "tilelang"
     (pkg / "lib").mkdir(parents=True)
@@ -1306,7 +1316,216 @@ def test_find_real_libcudart_finds_cu13_wheel_layout(tmp_path, monkeypatch):
     monkeypatch.setattr(ctypes.util, "find_library", lambda _n: None)
     _real_glob = glob.glob
     monkeypatch.setattr(
-        glob, "glob", lambda p, *a, **k: ([] if p.startswith("/usr") else _real_glob(p, *a, **k))
+        glob, "glob", lambda p, *a, **k: [] if p.startswith("/usr") else _real_glob(p, *a, **k)
     )
 
     assert perf._find_real_libcudart() == os.path.realpath(real)
+
+
+# ---------------------------------------------------------------------------
+# Blackwell fla GDN autotune restriction (fla #913 / #1000): on sm100/sm120 the
+# unrestricted prepare_wy_repr_bwd autotune space can select grad-miscomputing
+# configs (live B200 Qwen3.6-35B-A3B SFT: grad_norm ~1e8 from the first logged
+# step, loss flat or collapsing at every LR, while H200 trained healthily). The
+# worker restricts the space in-process to the B200-validated config, and fails
+# CLOSED (disables fla -> pure-PyTorch delta) when it cannot.
+# ---------------------------------------------------------------------------
+def _blackwell_torch(cc=(10, 0)):
+    """A stub ``torch`` that looks like a Blackwell card with CUDA available."""
+    t = types.ModuleType("torch")
+    t.cuda = types.SimpleNamespace(
+        is_available=lambda: True,
+        get_device_capability=lambda *a: cc,
+    )
+    return t
+
+
+class _TunerCfg:
+    def __init__(self, num_warps: int, num_stages: int) -> None:
+        self.num_warps = num_warps
+        self.num_stages = num_stages
+
+
+class _FakeTuner:
+    def __init__(self, configs) -> None:
+        self.configs = configs
+
+
+class _TunerWrapper:
+    """Decorator layer (heuristics/cache wrapper) the unwrap walk descends through via ``.fn``."""
+
+    def __init__(self, fn) -> None:
+        self.fn = fn
+
+
+_FULL_SPACE = [(2, 2), (2, 3), (2, 4), (4, 2), (4, 3), (4, 4)]
+
+
+def _patch_blackwell_stack(monkeypatch, *, cc=(10, 0), fla_present: bool = True, kernel=None):
+    """Wire the perf helper's touchpoints for the Blackwell autotune-restriction tests.
+
+    ``kernel`` is what ``fla.ops.gated_delta_rule.wy_fast.prepare_wy_repr_bwd_kernel``
+    resolves to (None = attribute absent). Returns (perf, removed, wy_fast_module).
+    """
+    import importlib.util
+
+    from flash.engine.worker import perf
+
+    monkeypatch.setitem(sys.modules, "torch", _blackwell_torch(cc))
+    monkeypatch.setattr(
+        importlib.util,
+        "find_spec",
+        (lambda name: object()) if fla_present else (lambda name: None),
+        raising=True,
+    )
+    wy = types.ModuleType("fla.ops.gated_delta_rule.wy_fast")
+    if kernel is not None:
+        wy.prepare_wy_repr_bwd_kernel = kernel
+    pkg_fla = types.ModuleType("fla")
+    pkg_ops = types.ModuleType("fla.ops")
+    pkg_gdr = types.ModuleType("fla.ops.gated_delta_rule")
+    pkg_fla.__path__, pkg_ops.__path__, pkg_gdr.__path__ = [], [], []
+    pkg_gdr.wy_fast = wy
+    for name, mod in (
+        ("fla", pkg_fla),
+        ("fla.ops", pkg_ops),
+        ("fla.ops.gated_delta_rule", pkg_gdr),
+        ("fla.ops.gated_delta_rule.wy_fast", wy),
+    ):
+        monkeypatch.setitem(sys.modules, name, mod)
+    removed: list[int] = []
+    monkeypatch.setattr(
+        perf, "_remove_fla_from_disk", lambda: (removed.append(1), (["/x/fla"], False))[1]
+    )
+    return perf, removed, wy
+
+
+def test_blackwell_gdn_autotune_restricted_to_validated_config(monkeypatch):
+    """sm100 + full upstream config space -> only num_warps=2/num_stages=4 survives; fla is KEPT.
+    The tuner is wrapped in decorator layers, so the unwrap walk (.fn descent) is exercised too."""
+    tuner = _FakeTuner([_TunerCfg(w, s) for w, s in _FULL_SPACE])
+    perf, removed, _ = _patch_blackwell_stack(
+        monkeypatch, kernel=_TunerWrapper(_TunerWrapper(tuner))
+    )
+
+    perf._restrict_fla_gdn_autotune_on_blackwell()
+
+    assert [(c.num_warps, c.num_stages) for c in tuner.configs] == [(2, 4)]
+    assert not removed, "restriction succeeded -> fla must be KEPT"
+
+
+def test_blackwell_gdn_autotune_sm120_also_restricted(monkeypatch):
+    """Consumer Blackwell (sm120) matches upstream's IS_NVIDIA_BLACKWELL gate: restricted too."""
+    tuner = _FakeTuner([_TunerCfg(w, s) for w, s in _FULL_SPACE])
+    perf, removed, _ = _patch_blackwell_stack(monkeypatch, cc=(12, 0), kernel=tuner)
+
+    perf._restrict_fla_gdn_autotune_on_blackwell()
+
+    assert [(c.num_warps, c.num_stages) for c in tuner.configs] == [(2, 4)]
+    assert not removed
+
+
+def test_blackwell_gdn_autotune_already_restricted_is_noop(monkeypatch):
+    """A pinned fla that already carries fla #1000 (space == validated config) is left alone —
+    the guard must be a no-op on a future pin bump, not a second filter or a fail-closed."""
+    tuner = _FakeTuner([_TunerCfg(2, 4)])
+    perf, removed, _ = _patch_blackwell_stack(monkeypatch, kernel=tuner)
+
+    perf._restrict_fla_gdn_autotune_on_blackwell()
+
+    assert [(c.num_warps, c.num_stages) for c in tuner.configs] == [(2, 4)]
+    assert not removed
+
+
+def test_blackwell_gdn_autotune_fail_closed_when_tuner_missing(monkeypatch):
+    """Autotuner not found (fla layout changed, e.g. an unreviewed pin bump) -> fail CLOSED:
+    fla is physically removed so transformers falls back to the pure-PyTorch delta rule."""
+    perf, removed, _ = _patch_blackwell_stack(monkeypatch, kernel=None)
+
+    perf._restrict_fla_gdn_autotune_on_blackwell()
+
+    assert removed, "no autotuner -> must DISABLE fla (grad-correctness over speed)"
+
+
+def test_blackwell_gdn_autotune_fail_closed_when_no_validated_config(monkeypatch):
+    """The validated (2,4) config missing from the space -> nothing safe to run -> fail CLOSED."""
+    tuner = _FakeTuner([_TunerCfg(4, 2), _TunerCfg(4, 3)])
+    perf, removed, _ = _patch_blackwell_stack(monkeypatch, kernel=tuner)
+
+    perf._restrict_fla_gdn_autotune_on_blackwell()
+
+    assert removed, "no validated config in space -> must DISABLE fla"
+
+
+def test_blackwell_gdn_autotune_noop_without_fla(monkeypatch):
+    """fla absent (already the pure-PyTorch path) -> nothing to restrict, nothing to remove."""
+    perf, removed, _ = _patch_blackwell_stack(monkeypatch, fla_present=False, kernel=None)
+
+    perf._restrict_fla_gdn_autotune_on_blackwell()
+
+    assert not removed
+
+
+def test_non_blackwell_gdn_autotune_untouched(monkeypatch):
+    """Hopper (sm90) keeps the full autotune space — the restriction is Blackwell-only (the
+    tilelang fast path owns the sm90 story; shrinking its Triton space would only cost perf)."""
+    tuner = _FakeTuner([_TunerCfg(w, s) for w, s in _FULL_SPACE])
+    perf, removed, _ = _patch_blackwell_stack(monkeypatch, cc=(9, 0), kernel=tuner)
+
+    perf._restrict_fla_gdn_autotune_on_blackwell()
+
+    assert len(tuner.configs) == len(_FULL_SPACE)
+    assert not removed
+
+
+# ---------------------------------------------------------------------------
+# sm100 tilelang GDN opt-out: the baked tilelang backend (needed for Hopper, fla
+# #640) computes WRONG GRADIENTS on B200/sm100 (measured: dq/dk ~0.72, dg ~1.28
+# rel-err at the production H==HV call shapes, deterministic, bf16 AND fp32) —
+# the root cause of the B200 35B-A3B SFT incident. The worker must opt fla out
+# via FLA_TILELANG=0 (upstream's own knob; upstream default-gates tilelang to
+# Hopper since fla #975) so fla dispatches to its Triton path, correct on sm100.
+# ---------------------------------------------------------------------------
+def _patch_arch(monkeypatch, cc):
+    from flash.engine.worker import perf
+
+    monkeypatch.setitem(sys.modules, "torch", _blackwell_torch(cc))
+    return perf
+
+
+def test_sm100_fla_tilelang_opted_out(monkeypatch):
+    """sm100 + FLA_TILELANG unset -> the worker sets FLA_TILELANG=0 before any fla dispatch."""
+    import os
+
+    perf = _patch_arch(monkeypatch, (10, 0))
+    monkeypatch.delenv("FLA_TILELANG", raising=False)
+
+    perf._force_fla_triton_gdn_on_sm100()
+
+    assert os.environ.get("FLA_TILELANG") == "0"
+
+
+def test_sm100_fla_tilelang_explicit_preset_wins(monkeypatch):
+    """An explicitly pre-set FLA_TILELANG (e.g. testing a fixed tilelang) is respected."""
+    import os
+
+    perf = _patch_arch(monkeypatch, (10, 0))
+    monkeypatch.setenv("FLA_TILELANG", "1")
+
+    perf._force_fla_triton_gdn_on_sm100()
+
+    assert os.environ.get("FLA_TILELANG") == "1"
+
+
+@pytest.mark.parametrize("cc", [(9, 0), (12, 0), (8, 9)])
+def test_non_sm100_fla_tilelang_untouched(monkeypatch, cc):
+    """sm90 NEEDS tilelang (fla #640); sm89/sm120 train healthily under the pin's default —
+    the opt-out is strictly sm100 (the measured-broken arch)."""
+    import os
+
+    perf = _patch_arch(monkeypatch, cc)
+    monkeypatch.delenv("FLA_TILELANG", raising=False)
+
+    perf._force_fla_triton_gdn_on_sm100()
+
+    assert "FLA_TILELANG" not in os.environ