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

{freesolo_flash_dev-0.2.41 → 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.41 → freesolo_flash_dev-0.2.43}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.41
+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.41 → freesolo_flash_dev-0.2.43}/flash/cli/__init__.py

@@ -249,7 +249,31 @@ def _build_parser() -> argparse.ArgumentParser:
         const="multi",
         help="scaffold a multi-turn environment (bounded episode with step_episode / score_episode).",
     )
-    setup.set_defaults(func=cmd_env_setup, turn_mode="single")
+    setup_reason = setup.add_mutually_exclusive_group()
+    setup_reason.add_argument(
+        "--reasoning",
+        dest="reasoning",
+        action="store_const",
+        const=True,
+        help="scaffold configs with reasoning enabled (thinking = true, with a raised token budget).",
+    )
+    setup_reason.add_argument(
+        "--no-reasoning",
+        dest="reasoning",
+        action="store_const",
+        const=False,
+        help="scaffold configs without reasoning. This is the default.",
+    )
+    setup.add_argument(
+        "-y",
+        "--yes",
+        dest="yes",
+        action="store_true",
+        help="accept defaults without prompting (single-turn, no reasoning).",
+    )
+    # turn_mode / reasoning default to None (unset) so the scaffold can tell an explicit flag
+    # apart from "not chosen yet" and, on an interactive terminal, ask instead of assuming.
+    setup.set_defaults(func=cmd_env_setup, turn_mode=None, reasoning=None, yes=False)
 
     env_list = env_sub.add_parser("list", help="list local environment sources")
     env_list.set_defaults(func=cmd_env_list)

{freesolo_flash_dev-0.2.41 → freesolo_flash_dev-0.2.43}/flash/cli/commands.py

@@ -330,14 +330,40 @@ _STARTER_DATASET_MULTITURN_JSONL = """\
 """
 
 
+_TURN_OPTIONS = [
+    ("single", "single-turn", "one prompt -> one response"),
+    ("multi", "multi-turn", "bounded episode: step_episode / score_episode"),
+]
+_REASONING_OPTIONS = [
+    ("off", "no reasoning", "the model answers directly"),
+    ("on", "reasoning", "thinking = true; spends more tokens per answer"),
+]
+
+
+def _setup_interactive(args) -> bool:
+    """Whether to ask the setup questions interactively.
+
+    Prompt only on a real terminal a human can answer from. Fall back to defaults (no prompt) when
+    --yes is passed, under CI, or when stdin is closed/redirected, so automation never blocks on an
+    unanswered prompt. A pseudo-TTY in CI can report isatty()=True, so CI is checked explicitly; a
+    closed fd 0 can make sys.stdin None, so that is guarded before .isatty()."""
+    if getattr(args, "yes", False):
+        return False
+    if os.environ.get("CI", "").strip().lower() not in ("", "0", "false", "no"):
+        return False
+    stdin = sys.stdin
+    if stdin is None or not stdin.isatty():
+        return False
+    return render.styled()
+
+
 def cmd_env_setup(args) -> int:
-    requested_multi = getattr(args, "turn_mode", "single") == "multi"
     starter_env = Path("environment.py")
     dataset = Path("dataset/train.jsonl")
     # An existing environment.py is the authoritative signal for which turn mode this
     # scaffold already uses (the dataset is plain JSONL with no reliable mode marker).
     # Anchor to it so a re-run never leaves a single-turn env beside a multi-turn
-    # dataset (or vice versa); the flag only decides the mode when starting fresh.
+    # dataset (or vice versa); the flag/answer only decides the mode when starting fresh.
     existing_multi: bool | None = None
     anchor = "environment.py"
     if starter_env.exists():
@@ -347,15 +373,59 @@ def cmd_env_setup(args) -> int:
         # distinctive prompt; use it so we don't drop a single-turn env beside it.
         existing_multi = "secret whole number" in dataset.read_text(encoding="utf-8")
         anchor = "dataset/train.jsonl"
-    if existing_multi is not None and existing_multi != requested_multi:
-        have = "multi-turn" if existing_multi else "single-turn"
-        want = "multi-turn" if requested_multi else "single-turn"
-        msg = (
-            f"existing {anchor} is {have}; keeping it and ignoring --{want}. "
-            f"Delete environment.py and dataset/train.jsonl first to re-scaffold as {want}."
+
+    # Resolve the turn mode. An existing scaffold wins (warn if a flag disagrees); otherwise an
+    # explicit --single-turn/--multi-turn flag wins; otherwise ask on a terminal, else single-turn.
+    flag_mode = getattr(args, "turn_mode", None)
+    if existing_multi is not None:
+        if flag_mode is not None and (flag_mode == "multi") != existing_multi:
+            have = "multi-turn" if existing_multi else "single-turn"
+            want = "multi-turn" if flag_mode == "multi" else "single-turn"
+            msg = (
+                f"existing {anchor} is {have}; keeping it and ignoring --{want}. "
+                f"Delete environment.py and dataset/train.jsonl first to re-scaffold as {want}."
+            )
+            print(render.warn(msg) if render.styled() else f"warning: {msg}", file=sys.stderr)
+        multi_turn = existing_multi
+    elif flag_mode is not None:
+        multi_turn = flag_mode == "multi"
+    elif _setup_interactive(args):
+        multi_turn = (
+            render.select("How does the model interact with your task?", _TURN_OPTIONS) == "multi"
         )
-        print(render.warn(msg) if render.styled() else f"warning: {msg}", file=sys.stderr)
-    multi_turn = requested_multi if existing_multi is None else existing_multi
+    else:
+        multi_turn = False
+
+    # Resolve reasoning. Like the turn mode, an existing config is authoritative: configs are only
+    # written when absent, so applying a reasoning flag to an already-scaffolded project would
+    # silently no-op (or write one config with reasoning and leave the other without). Anchor to the
+    # existing config's `thinking` state and warn if a flag disagrees; otherwise the flag wins;
+    # otherwise ask on a terminal; else off.
+    rl = Path("configs/rl.toml")
+    sft = Path("configs/sft.toml")
+    existing_reasoning: bool | None = None
+    for cfg in (rl, sft):
+        if cfg.exists():
+            existing_reasoning = "thinking = true" in cfg.read_text(encoding="utf-8")
+            break
+    flag_reason = getattr(args, "reasoning", None)
+    if existing_reasoning is not None:
+        if flag_reason is not None and flag_reason != existing_reasoning:
+            have = "reasoning" if existing_reasoning else "no reasoning"
+            want = "reasoning" if flag_reason else "no-reasoning"
+            msg = (
+                f"existing configs are {have}; keeping them and ignoring --{want}. "
+                f"Delete configs/rl.toml and configs/sft.toml first to re-scaffold with --{want}."
+            )
+            print(render.warn(msg) if render.styled() else f"warning: {msg}", file=sys.stderr)
+        reasoning = existing_reasoning
+    elif flag_reason is not None:
+        reasoning = flag_reason
+    elif _setup_interactive(args):
+        reasoning = render.select("Train with reasoning (thinking)?", _REASONING_OPTIONS) == "on"
+    else:
+        reasoning = False
+
     env_py = _STARTER_ENV_MULTITURN_PY if multi_turn else _STARTER_ENV_PY
     dataset_jsonl = _STARTER_DATASET_MULTITURN_JSONL if multi_turn else _STARTER_DATASET_JSONL
     Path("configs").mkdir(exist_ok=True)
@@ -373,28 +443,47 @@ def cmd_env_setup(args) -> int:
         'id = ""\n\n'
         '# secrets = ["SERPAPI_API_KEY"]\n\n'
     )
-    rl = Path("configs/rl.toml")
+    # `thinking = true` opts the run into reasoning mode. Reasoning shares the generation budget with
+    # the answer, so GRPO also gets a raised max_tokens. These strings are empty when reasoning is
+    # off, keeping the default scaffold byte-for-byte identical.
+    thinking_line = "thinking = true\n" if reasoning else ""
+    rl_reasoning_train = (
+        "max_tokens = 2048  # reasoning shares this budget with the answer; raised so it isn't truncated\n"
+        if reasoning
+        else ""
+    )
+    sft_reasoning_note = (
+        "# reasoning is on (thinking = true): each gold `output` must contain a <think>...</think>\n"
+        "# block; validate locally with freesolo.datasets.warn_missing_think_tags before a real run.\n"
+        if reasoning
+        else ""
+    )
     if not rl.exists():
         rl.write_text(
             'model = "Qwen/Qwen3.5-4B"\n'
-            'algorithm = "grpo"\n\n'
+            'algorithm = "grpo"\n'
+            f"{thinking_line}"
+            "\n"
             f"{env_comment}"
             "[train]\n"
             "steps = 150\n"
+            f"{rl_reasoning_train}"
             "lora_rank = 32\n"
             "# GPU and HF artifacts are managed automatically by the platform: the GPU is\n"
             "# the cheapest fitting managed class, and artifacts live in a private environment-scoped repo.\n"
         )
-    sft = Path("configs/sft.toml")
     if not sft.exists():
         sft.write_text(
             'model = "Qwen/Qwen3.5-4B"\n'
-            'algorithm = "sft"\n\n'
+            'algorithm = "sft"\n'
+            f"{thinking_line}"
+            "\n"
             f"{env_comment}"
             "[train]\n"
             "epochs = 1\n"
             "max_examples = 2  # rows to train on; the starter dataset has 2 (raise as your dataset grows)\n"
             "lora_rank = 32\n"
+            f"{sft_reasoning_note}"
             "# GPU and HF artifacts are managed automatically by the platform: the GPU is\n"
             "# the cheapest fitting managed class, and artifacts live in a private environment-scoped repo.\n"
         )
@@ -768,7 +857,9 @@ def cmd_deploy(args) -> int:
                 f"OpenAI-compatible base URL: {openai_base} — point clients at this /v1 base, "
                 "not the bare endpoint (which 404s on /chat/completions)."
             )
-            print(render.arrow(url_note) if render.styled() else f"note: {url_note}", file=sys.stderr)
+            print(
+                render.arrow(url_note) if render.styled() else f"note: {url_note}", file=sys.stderr
+            )
     if dep.get("state") != "dry_run":
         state = dep.get("state", "deploying")
         if state == "failed":

{freesolo_flash_dev-0.2.41 → freesolo_flash_dev-0.2.43}/flash/cli/render.py

@@ -273,6 +273,38 @@ def arrow(msg: str) -> str:
     return _safe(f"{_paint(_glyph('→', '->'), _ACCENT2)} {_dim(msg)}")
 
 
+def select(title: str, options: list[tuple[str, str, str]], default: int = 0) -> str:
+    """Themed single-choice prompt; returns the chosen option's value.
+
+    ``options`` is a list of ``(value, label, hint)``. The default option is marked and taken on
+    an empty answer (enter). Reads via ``input()`` so it is easy to drive in tests; the caller
+    decides *when* to prompt (interactive stdin), so this always prompts when called. On EOF the
+    default is returned so a closed stdin never hangs the scaffold."""
+    q = _paint(_glyph("?", "?"), _ACCENT, "1")
+    print(f"{q} {_bold(_safe(title))}")
+    for i, (_value, label, hint) in enumerate(options):
+        num = _paint(f"{i + 1})", _ACCENT2)
+        lab = _bold(label) if i == default else label
+        tail = f"  {_dim(_safe(hint))}" if hint else ""
+        mark = _paint(" (default)", _GREEN) if i == default else ""
+        print(f"  {num} {_safe(lab)}{tail}{mark}")
+    pointer = _paint(_glyph("›", ">"), _ACCENT2)  # noqa: RUF001 (the glyph is the point)
+    for _ in range(5):
+        try:
+            raw = input(f"{pointer} ").strip()
+        except EOFError:
+            print()
+            return options[default][0]
+        if not raw:
+            return options[default][0]
+        if raw.isdigit():
+            idx = int(raw) - 1
+            if 0 <= idx < len(options):
+                return options[idx][0]
+        print(note(f"enter 1-{len(options)}, or press enter for the default"))
+    return options[default][0]
+
+
 def money(value: float, decimals: int = 4) -> str:
     return _paint(f"${value:.{decimals}f}", _TEAL)
 

{freesolo_flash_dev-0.2.41 → 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.41 → 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.41 → 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.41 → 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.41 → 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",
@@ -110,12 +112,13 @@ def _effective_worker_env(spec=None) -> dict[str, str]:
 # Chalk is published to PUBLIC PyPI on every version bump (chalk .github/workflows/publish.yml).
 # Pin the exact PUBLIC version so worker + kernel-cache-bake pods install with NO GitHub auth — the
 # chalk repo is INTERNAL, so a git+https default made installs fail wherever GITHUB_TOKEN is absent
-# (bake pod, tokenless workers). freesolo-chalk 0.5.0 = chalk PR #21 fafe167 (standalone, Liger fully
-# replaced, + GRPO op + LoRA grad-gate). Bump DEFAULT_CHALK_VERSION when the merged kernel surface moves.
-DEFAULT_CHALK_VERSION = "0.5.1"
+# (bake pod, tokenless workers). freesolo-chalk 0.5.2 = chalk dev->main sync #44 (standalone, Liger
+# fully replaced; adds arch-TUNED gdn_gated_rmsnorm + lora on sm89/sm90/sm100, MoE grouped-GEMM,
+# serving surface, MiniCPM/Llama kernels). Bump DEFAULT_CHALK_VERSION when the merged kernel surface moves.
+DEFAULT_CHALK_VERSION = "0.5.2"
 DEFAULT_CHALK_SPEC = f"freesolo-chalk=={DEFAULT_CHALK_VERSION}"
 # Provenance only (kernel-cache fingerprint / traceability): the chalk commit this version was cut from.
-LATEST_CHALK_MAIN_SHA = "fafe167158eb3b7c947e7182ea4eddc351e4ca7d"
+LATEST_CHALK_MAIN_SHA = "37b2fb9e1a123ceb902ba6ffc4ed182182c6e7d7"
 
 
 def chalk_extra_pip(spec=None) -> list[str]:

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

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "0.2.41"
+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.41"
+version = "0.2.43"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]