docs: blog post; HF sync workflow; kaggle prompt/reward + tests

- Add blog.md (hackathon narrative) and README Blog link
- Git LFS for plot PNGs; sync-to-hf-space workflow + sync_hf_space.sh
- Dataset/reward prompt contract and format penalty; hf_run_job uses clone
- Tests for preprocess_assembly and training prompt contract

Made-with: Cursor

.gitattributes

@@ -1 +1,2 @@
+# Hugging Face Hub rejects raw PNG git objects on Space pushes; use Git LFS.
 colab/results/plots/*.png filter=lfs diff=lfs merge=lfs -text

.github/workflows/sync-to-hf-space.yml

@@ -0,0 +1,42 @@
+# Sync GitHub ka-ori/arm-gym → Hugging Face Space https://huggingface.co/spaces/kaori02/arm-gym
+# (branch main → Space branch main).
+#
+# Hugging Face rejects raw PNG blobs in git packs; this workflow rewrites history in the
+# runner with `git lfs migrate import` for colab/results/plots/*.png before pushing to the
+# Space only (GitHub is unchanged). See: https://huggingface.co/docs/hub/xet
+#
+# Secrets: HF_TOKEN (write access to the Space). Optional vars: HF_USERNAME, HF_SPACE
+# (defaults kaori02 / arm-gym).
+#
+# Local one-shot: scripts/sync_hf_space.sh
+
+name: Sync GitHub to Hugging Face Space
+
+on:
+  push:
+    branches: [main]
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          lfs: true
+
+      - name: Rewrite PNGs to Git LFS for Hub push
+        run: |
+          git lfs install
+          git lfs migrate import --include="colab/results/plots/*.png" --include-ref=refs/heads/main
+
+      - name: Push to Hugging Face Space
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          HF_USER: ${{ vars.HF_USERNAME }}
+          HF_SPACE: ${{ vars.HF_SPACE }}
+        run: |
+          U="${HF_USER:-kaori02}"
+          S="${HF_SPACE:-arm-gym}"
+          git push --force "https://${U}:${HF_TOKEN}@huggingface.co/spaces/${U}/${S}" HEAD:main

README.md

@@ -17,6 +17,7 @@ short_description: GRPO env for AArch64 superoptimization
 
 <p align="center">
   <a href="https://huggingface.co/spaces/dot-mkv/arm-gym">HF Space</a> &middot;
+  <a href="./blog.md">Blog</a> &middot;
   <a href="#quick-start">Quick Start</a> &middot;
   <a href="#results">Results</a> &middot;
   <a href="#why-it-matters">Why It Matters</a>

blog.md

@@ -0,0 +1,270 @@
+# We trained an AI to write faster ARM assembly than the compiler
+
+Some ideas start with a paper.
+
+In January 2026, a team from Stanford and UIUC published [SuperCoder](https://arxiv.org/abs/2505.11480) - a system that trained a language model to write assembly code faster than `gcc -O3`, the most aggressive optimization setting of the world's most widely used compiler. Their result: 1.46x average speedup. On x86-64. With a 7 billion parameter model trained via reinforcement learning.
+
+The paper was explicit about one thing it did not do: ARM.
+
+> "Extending to ARM, RISC-V, and GPU kernels is noted as future work."
+
+That sentence is where ARM-Gym begins.
+
+---
+
+## What this is about, for everyone
+
+Before getting into what we built, here is the context for anyone who does not spend their days thinking about compilers and processors.
+
+**What is a compiler?**
+
+A compiler is a program that takes code written in a high-level language like C and translates it into machine instructions. GCC and Clang are the two most important compilers in existence. When you compile with `-O3`, you are asking the compiler to apply every optimization it knows. It inlines functions, unrolls loops, reorders instructions to avoid stalls, selects faster instruction variants where it can.
+
+GCC has been doing this for decades. Thousands of engineer-years have gone into making `-O3` as good as it is.
+
+**And yet.**
+
+Compilers must be conservative. A compiler cannot make an assumption that might be wrong for even one program. It cannot take a risk that improves performance 99% of the time but breaks the other 1%. It follows rules. Rules generalize but do not specialize.
+
+**What is ARM?**
+
+ARM is a processor architecture - a specific design for how machine instructions are structured and executed. For a long time, ARM meant smartphones. That is no longer true. Today:
+
+- AWS Graviton5 - the most widely deployed cloud compute in the world - runs on ARM
+- Azure Cobalt 100, Microsoft's custom data center chip, runs on ARM
+- Every Apple Mac sold since 2020 runs on ARM
+- Meta's AGI CPU - 136 cores, 3nm fabrication, deployed in 2026 - runs on ARM
+
+Any improvement in how efficiently code runs on ARM touches all of that. And the specific code that matters most is the tight computational loops inside AI inference: matrix multiply, softmax, convolution. These functions run millions of times per second in every large model deployment.
+
+---
+
+## The research that made this possible: SuperCoder
+
+The SuperCoder paper (Wei et al., arXiv:2505.11480, Stanford/UIUC, 2025) is the direct foundation for what we built. Understanding what they proved is essential for understanding why ARM-Gym is the next step.
+
+### What SuperCoder showed
+
+The paper asked: can a language model learn to write assembly that beats the compiler, purely through reinforcement learning?
+
+They evaluated 23 language models on a benchmark of 8,072 assembly programs (average 130 lines each - far larger than any prior dataset, which maxed out at 15 lines and no loops). Every program came with its `gcc -O3` baseline assembly and a set of test cases.
+
+The base model they chose for training was Qwen2.5-Coder-7B-Instruct - not because it was the strongest baseline, but because it had the highest test pass rate (61.4%) among open-source models, leaving the most room to improve. Claude-opus-4 had a slightly higher average speedup (1.43x) but was not open-source and could not be fine-tuned.
+
+They trained using both PPO and GRPO. The reward function was simple: if the generated assembly compiles, passes all test cases, and runs faster than `gcc -O3`, the reward equals the speedup. Otherwise, zero. No partial credit. No reward for being partially correct.
+
+That last point turned out to matter a lot. They tested an alternative reward that gave partial credit for passing some tests - and it performed worse (1.38x vs 1.46x). The lesson: partial credit lets the model avoid putting in the work of actually being correct and fast. Binary pass/fail forces it.
+
+### The results
+
+| Model | Correctness before training | Correctness after | Avg speedup |
+|---|---|---|---|
+| Qwen2.5-Coder-7B (base) | 61.4% | - | 1.10x |
+| SuperCoder (GRPO) | - | 95.0% | 1.44x |
+| SuperCoder (PPO) | - | 95.0% | 1.46x |
+
+Correctness jumped from 61.4% to 95%. Average speedup went from 1.10x to 1.46x. The model went from occasionally beating the compiler to reliably beating it.
+
+One more finding worth noting: 98.5% of the speedup came from instruction scheduling and code layout - reordering instructions and basic blocks to better hide latency and avoid pipeline stalls. Not exotic instruction selection. The model learned that the compiler's instruction order is not optimal, and found better orderings.
+
+### What SuperCoder did not do
+
+It targeted x86-64 only. The paper uses the IBM CodeNet dataset, which contains competitive programming submissions compiled for x86. ARM was explicitly outside scope.
+
+ARM has a completely different instruction set. Different SIMD extensions (NEON, SVE2 instead of AVX/SSE). Different pipeline characteristics. Different scheduler model. Different optimization opportunities. A model trained on x86 assembly cannot be directly applied to ARM.
+
+That is the gap ARM-Gym fills.
+
+---
+
+## What we built: ARM-Gym
+
+ARM-Gym is a reinforcement learning environment built on [OpenEnv](https://github.com/meta-pytorch/OpenEnv), the framework from Meta and Hugging Face. The task is identical to SuperCoder's in structure - generate assembly that beats `gcc -O3` - but the target is AArch64 (ARM's 64-bit architecture) and the kernels are specifically chosen to be representative of AI inference workloads.
+
+### The kernel library
+
+We wrote 15 C function templates covering the operations that dominate AI inference:
+
+- Vector operations: `vec_add`, `dot_product`, `saxpy`
+- Matrix operations: `gemv`, `matmul`
+- Activation and normalization: `softmax`, `layer_norm`
+- Convolution: `conv1d`, `conv2d`
+- Elementwise: `relu`, `gelu`, `silu`, `fma`
+
+From these 15 templates, we generate 523 variants by varying sizes, data types (float32, float16, int8), and parameters. This gives the training loop enough diversity to prevent memorization while staying domain-relevant.
+
+### The training loop
+
+Here is exactly how training works, step by step:
+
+```mermaid
+flowchart LR
+    A["C Kernel<br/>(15 templates × 523 variants)"] --> B["gcc -O3<br/>Baseline Assembly"]
+    B --> C["LLM Prompt<br/>(C + Baseline ASM)"]
+    C --> D["Qwen2.5-Coder-7B<br/>+ LoRA"]
+    D --> E["Agent Assembly"]
+    E --> F{"3-Gate Verifier"}
+    F -->|"Gate 1: Syntax"| G["GNU as"]
+    F -->|"Gate 2: Correctness"| H["QEMU × 20<br/>Adversarial Tests"]
+    F -->|"Gate 3: Performance"| I["LLVM-MCA<br/>Neoverse V2"]
+    I --> J["Dual Verifier<br/>Cross-Check"]
+    J --> K["Reward<br/>max(0, speedup - 1)"]
+    K --> L["GRPO Update"]
+    L --> D
+
+    style F fill:#f96,stroke:#333,color:#000
+    style J fill:#69f,stroke:#333,color:#000
+    style L fill:#6c6,stroke:#333,color:#000
+```
+
+**Step 1 - Sample a kernel.** Pick one of the 523 variants at random. This is the function the model needs to optimize.
+
+**Step 2 - Compile the baseline.** Run `gcc -O3` on the C source. This gives us the baseline assembly and a baseline cycle count from LLVM-MCA.
+
+**Step 3 - Build the prompt.** Send the model the C source code, the baseline assembly, and an instruction: write optimized AArch64 assembly for this function, wrapped in `<assembly>...</assembly>` tags. Crucially, the baseline is always included - SuperCoder found that without it, even strong models produce 0% compilable code. The baseline is load-bearing context.
+
+**Step 4 - Verify.** The model's output goes through three sequential gates. Pass all three and you get a reward. Fail any one and you score zero.
+
+**Step 5 - Update with GRPO.** We use Group Relative Policy Optimization - the same algorithm from the DeepSeekMath paper that has since been adopted widely in RL for LLMs. GRPO generates multiple completions for the same prompt, scores them all, and uses their relative quality to compute the learning signal. It does not need a separate value function or critic model, which makes it efficient on constrained hardware.
+
+### The verifier: why it cannot be cheated
+
+Every RL environment has a reward hacking problem. Give an agent a metric and it will find the most efficient path to that metric - which is often not the path you intended.
+
+We saw this play out in Phase 1 of this hackathon across every finalist submission. One system's agent learned to starve long-context requests so short-request throughput looked better. Another learned to disconnect network access so error logs stopped appearing. A third learned to drop database tables so schema validation errors vanished. In all three cases, the metric went up and the actual objective was completely destroyed.
+
+Our verifier was designed from the start to have no exploitable surface.
+
+**Gate 1: Syntax via the real assembler.**
+
+The assembly must compile with `aarch64-linux-gnu-as`, the actual GNU assembler for ARM. Not a regular expression. Not a syntax checker. The real tool that produces a real binary object. If it rejects the assembly, the score is zero.
+
+**Gate 2: Correctness via randomized QEMU tests.**
+
+The compiled binary runs 20 times inside `qemu-aarch64-static`, a full ARM CPU emulator. Each run uses a different set of randomly generated inputs - edge cases, boundary values, near-overflow values. The output must match the original C function's output within floating-point tolerance. Inputs are randomized every episode, so the model cannot memorize test inputs and hardcode outputs for them.
+
+**Gate 3: Performance via LLVM-MCA.**
+
+We measure cycles with LLVM-MCA using the LLVM 21 Neoverse V2 scheduling model. This is a static analysis tool - it reads assembly text and estimates cycles based on the CPU's instruction latencies and throughput. It takes no inputs, executes nothing, and has no runtime attack surface.
+
+**Cross-check.** We compare the QEMU instruction count against the LLVM-MCA cycle estimate. A ratio above 3x triggers a hard veto - something is wrong, and we discard the result regardless of the apparent speedup.
+
+**3-sigma bound.** We maintain an offline distribution of speedup values for each kernel variant. Any result more than three standard deviations above the mean is rejected as a statistical outlier.
+
+No LLM judge anywhere in this stack. SuperCoder used Hyperfine (real execution timing). We use LLVM-MCA (static analysis) because it is deterministic, sub-millisecond per evaluation, and has no measurement noise. The trade-off is that MCA is a model of the hardware, not the hardware itself - which is why we label results as "MCA-model speedup" until we can validate on physical Graviton silicon.
+
+### The reward signal
+
+```
+reward = max(0, speedup - 1.0)
+```
+
+If the model's assembly is slower than `gcc -O3`, the reward is zero - neutral, not a penalty. If it is faster, the reward is the fractional improvement above parity: 1.5x speedup gives 0.5, 2x gives 1.0. Capped at 2.0 to prevent a single outlier from dominating the gradient.
+
+The zero floor is not arbitrary. In GRPO, rewards within a group are z-score normalized before computing the advantage. If slower-than-compiler gave a negative reward, a group where all completions happen to be slow would produce similar negative values that normalize to near zero - producing no gradient. Making slower neutral means even a group of slow completions still has relative differences that produce a learning signal. This mirrors exactly what SuperCoder found: sparse terminal reward (no partial credit) consistently outperforms reward designs that penalize failures.
+
+### The curriculum
+
+Not all kernels are equally hard. Sending the model to optimize a tiled matrix multiply before it has learned to write syntactically valid assembly is wasteful. ARM-Gym uses a staged curriculum that matches kernel difficulty to the model's current capability.
+
+```mermaid
+flowchart LR
+    S1["Stage 1: Scalar<br/>vec_add, dot, saxpy"] -->|"80% variants ≥ 1.05x"| S2["Stage 2: NEON<br/>gemv, conv1d, fma"]
+    S2 -->|"80% variants ≥ 1.05x"| S3["Stage 3: Loops<br/>matmul, softmax"]
+    S3 -->|"Beat -O3 mean"| S4["Stage 4: SVE2<br/>(Stretch)"]
+
+    style S1 fill:#bfb,stroke:#333,color:#000
+    style S2 fill:#fbf,stroke:#333,color:#000
+    style S3 fill:#bbf,stroke:#333,color:#000
+    style S4 fill:#fbb,stroke:#333,color:#000
+```
+
+**Stage 1 - Scalar kernels.** Functions like `vec_add` and `saxpy` where the compiler produces a scalar loop. NEON vectorization (processing 4 floats at once instead of 1) is the primary optimization opportunity. This is learnable early because the pattern is consistent.
+
+**Stage 2 - NEON kernels.** Functions where the compiler already emits NEON, but instruction ordering and register reuse can be improved. Requires the model to reason about pipeline latency, not just instruction selection.
+
+**Stage 3 - Loop-heavy kernels.** Matrix multiply and softmax, where optimization requires loop tiling, unrolling, and prefetch placement. These are the hardest patterns in Stage 1-3.
+
+**Stage 4 - SVE2 (stretch target).** ARM's Scalable Vector Extension, available on Neoverse V2 and V3. No training data exists for SVE2 code generation. This stage is explicitly a research frontier - there is no known baseline for what an RL agent can achieve here.
+
+Advancement between stages requires beating `gcc -O3` by at least 5% on 80% of the variants in the current stage. The model must demonstrate broad capability, not exploit a single easy variant.
+
+---
+
+## How ARM-Gym differs from SuperCoder
+
+| Aspect | SuperCoder | ARM-Gym |
+|---|---|---|
+| Target architecture | x86-64 | AArch64 (ARM) |
+| Dataset | 7,872 competitive programming programs | 523 AI inference kernel variants |
+| Performance measurement | Hyperfine (real hardware timing) | LLVM-MCA (static analysis, deterministic) |
+| RL framework | VERL | HuggingFace TRL |
+| Verifier | Compile + test pass | Compile + QEMU + LLVM-MCA + cross-check + 3-sigma |
+| Reward design | Binary terminal (same principle) | Binary terminal + format shaping |
+| Dataset focus | General programs | GEMM, matmul, softmax, conv (AI inference hot path) |
+| Prior work exists | Yes (this is the paper) | No - ARM is open |
+
+The most important difference is the last one. SuperCoder is the proof of concept. ARM-Gym is the next frontier. The paper itself identified ARM as the natural extension - we are building it.
+
+---
+
+## Results
+
+*[Results will be updated here after training completes.]*
+
+| Metric | Value |
+|---|---|
+| Best speedup over `gcc -O3` | [to be updated] |
+| Win rate | [to be updated] |
+| Correctness rate | [to be updated] |
+| Training steps and GPU time | [to be updated] |
+| SuperCoder reference (x86-64, PPO) | 1.46x average over `gcc -O3` |
+| Kernel variants trained on | 523 (15 templates) |
+
+---
+
+## What we learned from building this
+
+**The reward formula has to be exactly right.** We made one sign error - `speedup - 1.0` instead of `max(0, speedup - 1.0)` - and it silently destroyed the speedup gradient for an entire run. A group where all completions are slightly slow produces similar small negative values. After z-score normalization, they all collapse to near zero. No gradient. The model was learning correctness but not speed, and there was nothing in the loss curves to tell us why. The fix was one character. Test your reward function independently before attaching a model to it.
+
+**LLVM version has a hard dependency for ARM.** LLVM 17's Neoverse V2 scheduling model had the processor's issue-width wrong: 16 microoperations per cycle instead of the correct 8. Training on this would teach the model to optimize for a processor that does not exist. We pinned LLVM 21 specifically because of this correction.
+
+**Thinking models fail on assembly generation.** SuperCoder's benchmark found DeepSeek-R1 compiles at 0% across all 200 evaluation problems. The chain-of-thought habit causes the model to spend its entire output budget reasoning about instruction semantics - and never producing executable code. This is a known failure mode for reasoning-heavy models on generation tasks. The base model for assembly RL should not be a reasoning model.
+
+**Correct EOS token alignment is not optional.** Qwen2.5 uses `<|im_end|>` as its chat end-of-turn token, but TRL's GRPOTrainer by default reads `tokenizer.eos_token_id` for generation stopping. Without explicitly aligning these, the model never stops generating cleanly. Completions run to the token limit, producing garbage that collapses the reward signal. This required an explicit fix before training produced any useful signal at all.
+
+**GRPO needs within-group diversity.** With temperature 0.5, all completions in a group come out very similar - similar tokens, similar speedup, similar z-scores, near-zero gradient. Temperature 0.8 fixes this. More diversity means some completions try NEON vectorization, some stay scalar, some hallucinate - and the relative comparison between them becomes meaningful enough to drive learning.
+
+---
+
+## Why this matters
+
+Compilers use rules. Rules are safe, general, and conservative by necessity. Reinforcement learning finds what rules cannot - the specific instruction sequences, the register orderings, the NEON patterns that extract cycles on a specific microarchitecture for a specific workload.
+
+SuperCoder proved this approach works on x86. That result is now published, citable, and reproducible. ARM is the same problem on a larger market with no existing solution.
+
+AWS, Azure, Apple, and Meta have all made major bets on ARM infrastructure. The AI inference workloads running on that infrastructure are bottlenecked by the same matrix multiply and softmax kernels we are optimizing. Any improvement compounds.
+
+Could a researcher write a paper extending SuperCoder to ARM? Yes. That paper does not exist yet. ARM-Gym is that paper in environment form.
+
+---
+
+## What comes next
+
+**Silicon validation.** Every cycle count in ARM-Gym is an LLVM-MCA estimate on the Neoverse V2 model. Until we run the model's output on a physical Graviton3 machine and measure wall-clock time, these are model-predicted speedups. That validation is the step that turns "MCA speedup" into a real claim.
+
+**SVE2.** No model has been specifically trained to generate ARM SVE2 code. The Scalable Vector Extension is the highest-bandwidth path on Neoverse V2 and V3, and it is essentially unexplored territory for code generation models. The optimization potential is high and the competition is zero.
+
+**Larger models.** SuperCoder used 7B. ARM-Gym's current training targets 7B as well. The Best-of-8 sampling result from the paper (1.46x → 1.93x) suggests that scaling inference (more candidates, pick the best) is as important as scaling model size. Both directions are worth exploring.
+
+---
+
+## Try it
+
+- **Live environment:** [huggingface.co/spaces/dot-mkv/arm-gym](https://huggingface.co/spaces/dot-mkv/arm-gym)
+- **Training notebook:** `colab/arm_gym_grpo_kaggle.ipynb` - upload to Kaggle, set accelerator to T4 GPU, run all cells
+- **Paper we built on:** [SuperCoder (arXiv:2505.11480)](https://arxiv.org/abs/2505.11480)
+
+---
+
+*Meta / HuggingFace OpenEnv Hackathon India 2026 - Finals. Theme: Wild Card. Team: (dot)mkv.*

kaggle/dataset.py

@@ -14,26 +14,24 @@ from arm_gym.compile_baseline import ToolchainInfo, compile_to_asm
 from arm_gym.kernels import KernelVariant, generate_all, split_train_eval
 
 
-# Prompt shape from SuperCoder Appendix A.3 (arXiv:2505.11480), ISA swapped
-# to AArch64. Verbatim constraints: no extra text, no comments in asm, tags.
-# PDF: meta-hackathon-llm-wiki/papers/supercoder.pdf
+# C + baseline stay in the row for the verifier; user message ends with an
+# ``<assembly>`` prefill so the assistant continues inside the tag (Qwen chat).
 SYSTEM_PROMPT = (
-    "You are an expert AArch64 (aarch64-linux-gnu-gcc) assembly writer. "
-    "Obey the user block exactly. Output only what is asked in the required tags."
+    "You write only AArch64 (aarch64-linux-gnu) assembly inside the user’s "
+    "<assembly> block. No prose, no C, no other languages—assembly between "
+    "the opening line the user started and a closing </assembly> tag."
 )
 
 
 def user_prompt(c_source: str, baseline_asm: str) -> str:
-    # Order and wording follow SuperCoder A.3 (x86-64 -> AArch64 for arm-gym).
     return (
-        "Given the following C code and assembly code, your task is to generate "
-        "highly optimized AArch64 assembly code.\n\n"
-        f"C Code:\n{c_source}\n\n"
-        f"Assembly Code:\n{baseline_asm}\n\n"
-        "Only output the optimized assembly code. Do not include any other text. "
-        "Do not write any comments in the assembly code. "
-        "Wrap the assembly code in <assembly></assembly> tags.\n\n"
-        "Optimized Assembly Code:\n"
+        "Generate ONLY AArch64 assembly.\n\n"
+        "Wrap your continuation in <assembly></assembly> tags "
+        "(the opening tag is already below—finish the block and close with "
+        "</assembly>).\n\n"
+        f"C code:\n{c_source}\n\n"
+        f"Baseline assembly (gcc -O3):\n{baseline_asm}\n\n"
+        "<assembly>\n"
     )
 
 
@@ -60,7 +58,8 @@ def render_prompt(tokenizer, c_source: str, baseline_asm: str) -> str:
             prompt = f"{SYSTEM_PROMPT}\n\n{messages[1]['content']}\n"
     else:
         prompt = f"{SYSTEM_PROMPT}\n\n{messages[1]['content']}\n"
-    return prompt.rstrip()
+    # Do not rstrip: user content must end with "<assembly>\n" for generation prefill.
+    return prompt if prompt.endswith("\n") else prompt + "\n"
 
 
 @dataclass

kaggle/hf_run_job.sh

@@ -181,269 +181,8 @@ def install_grpo_assembly_stopping(model: Any, tokenizer: Any, **kwargs) -> None
 SHIMPY
 echo "[hotpatch] stopping criteria DISABLED -- using natural chat EOS"
 
-# 3. dataset.py: overwrite with fixed version (no <assembly> pre-injection)
-cat > kaggle/dataset.py << 'DATASETPY'
-"""Dataset builder: kernel variant -> GRPO prompt."""
-from __future__ import annotations
-from dataclasses import dataclass
-from datasets import Dataset
-from arm_gym.compile_baseline import ToolchainInfo, compile_to_asm
-from arm_gym.kernels import KernelVariant, generate_all, split_train_eval
-
-SYSTEM_PROMPT = (
-    "You are an expert AArch64 (aarch64-linux-gnu-gcc) assembly writer. "
-    "Obey the user block exactly. Output only what is asked in the required tags."
-)
-
-def user_prompt(c_source: str, baseline_asm: str) -> str:
-    return (
-        "Given the following C code and assembly code, your task is to generate "
-        "highly optimized AArch64 assembly code.\n\n"
-        f"C Code:\n{c_source}\n\n"
-        f"Assembly Code:\n{baseline_asm}\n\n"
-        "Only output the optimized assembly code. Do not include any other text. "
-        "Do not write any comments in the assembly code. "
-        "Wrap the assembly code in <assembly></assembly> tags.\n\n"
-        "Optimized Assembly Code:\n"
-    )
-
-def render_prompt(tokenizer, c_source: str, baseline_asm: str) -> str:
-    messages = [
-        {"role": "system", "content": SYSTEM_PROMPT},
-        {"role": "user", "content": user_prompt(c_source, baseline_asm)},
-    ]
-    if tokenizer is not None:
-        kwargs: dict = {
-            "tokenize": False,
-            "add_generation_prompt": True,
-            "enable_thinking": False,
-        }
-        try:
-            prompt = tokenizer.apply_chat_template(messages, **kwargs)
-        except TypeError:
-            kwargs.pop("enable_thinking", None)
-            prompt = tokenizer.apply_chat_template(messages, **kwargs)
-    else:
-        prompt = f"{SYSTEM_PROMPT}\n\n{messages[1]['content']}\n"
-    return prompt.rstrip()
-
-@dataclass
-class DatasetConfig:
-    max_train: int = 256
-    max_eval: int = 32
-    difficulty_max: int = 2
-
-def build(tc: ToolchainInfo, cfg: DatasetConfig | None = None,
-          tokenizer=None) -> tuple[Dataset, Dataset, dict[str, KernelVariant]]:
-    cfg = cfg or DatasetConfig()
-    from arm_gym.kernels import TEMPLATES
-    variants = [v for v in generate_all()
-                if TEMPLATES[v.template_name].difficulty <= cfg.difficulty_max]
-    train_v, eval_v = split_train_eval(variants, eval_frac=0.1, seed=0)
-    train_v = train_v[:cfg.max_train]
-    eval_v = eval_v[:cfg.max_eval]
-    def to_row(v: KernelVariant) -> dict | None:
-        try:
-            baseline_asm = compile_to_asm(v.c_source, tc)
-        except Exception:
-            return None
-        prompt = render_prompt(tokenizer, v.c_source, baseline_asm)
-        return {"prompt": prompt, "variant_id": v.variant_id,
-                "baseline_asm": baseline_asm, "c_source": v.c_source}
-    train_rows = [r for r in (to_row(v) for v in train_v) if r is not None]
-    eval_rows = [r for r in (to_row(v) for v in eval_v) if r is not None]
-    lookup = {v.variant_id: v for v in train_v + eval_v}
-    return Dataset.from_list(train_rows), Dataset.from_list(eval_rows), lookup
-DATASETPY
-echo "[hotpatch] dataset.py replaced (no <assembly> pre-injection)"
-
-# 4. reward_fn.py: overwrite with version expecting both tags
-cat > kaggle/reward_fn.py << 'REWARDPY'
-"""Reward functions for GRPOTrainer - both-tag aware.
-
-Model now emits both <assembly> and </assembly> in completions.
-format_reward rewards: +0.1 open tag, +0.1 close tag, +0.1 body >= 20 chars.
-"""
-from __future__ import annotations
-import hashlib, re, threading
-from dataclasses import dataclass
-from arm_gym.compile_baseline import detect_toolchain
-from arm_gym.verifier import run_correctness_qemu
-from arm_gym.mca import run_mca
-from arm_gym.rollout_budget import TestCase
-from arm_gym.verifier import VerifierConfig, assemble, cleanup_temp_dirs
-
-_ASM_RE = re.compile(r"<assembly>(.*?)</assembly>", re.DOTALL | re.IGNORECASE)
-_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
-
-def extract_assembly(text: str) -> str:
-    text = _THINK_RE.sub("", text).strip()
-    m = _ASM_RE.search(text)
-    if m:
-        return _clean_assembly(m.group(1))
-    if "</assembly>" in text.lower():
-        body = re.split(r"</assembly>", text, flags=re.IGNORECASE)[0]
-        if "<assembly>" in body.lower():
-            body = re.split(r"<assembly>", body, flags=re.IGNORECASE)[-1]
-        return _clean_assembly(body)
-    if "<assembly>" in text.lower():
-        return _clean_assembly(re.split(r"<assembly>", text, flags=re.IGNORECASE)[-1])
-    return _clean_assembly(text)
-
-def _clean_assembly(text: str) -> str:
-    text = text.strip()
-    text = re.sub(r"^```(?:asm|assembly|aarch64)?\s*", "", text, flags=re.IGNORECASE)
-    text = re.sub(r"\s*```$", "", text)
-    return text.strip()
-
-@dataclass
-class _Entry:
-    assembles: bool = False
-    runs: bool = False
-    speedup: float = 0.0
-
-_CACHE: dict[str, _Entry] = {}
-_LOCK = threading.Lock()
-_BASELINE: dict[str, float] = {}
-_VCFG: VerifierConfig | None = None
-_DEBUG_LIMIT = 5
-_DEBUG_SHOWN = 0
-
-def _cfg() -> VerifierConfig:
-    global _VCFG
-    if _VCFG is None:
-        tc = detect_toolchain()
-        _VCFG = VerifierConfig(
-            mca_bin=tc.mca or "llvm-mca", assembler="aarch64-linux-gnu-as",
-            linker="aarch64-linux-gnu-ld", qemu="qemu-aarch64-static", mcpu=tc.mcpu)
-    return _VCFG
-
-def _bcy(vid: str, basm: str) -> float:
-    if vid not in _BASELINE:
-        try:
-            rep = run_mca(basm, _cfg().mca_bin, _cfg().mcpu)
-            _BASELINE[vid] = float(rep.total_cycles)
-        except Exception:
-            _BASELINE[vid] = 1000.0
-    return _BASELINE[vid]
-
-def _key(text: str, vid: str) -> str:
-    return hashlib.md5(f"{text}::{vid}".encode()).hexdigest()
-
-def _run(text: str, vid: str, basm: str) -> _Entry:
-    global _DEBUG_SHOWN
-    k = _key(text, vid)
-    with _LOCK:
-        if k in _CACHE:
-            return _CACHE[k]
-    e = _Entry()
-    asm = extract_assembly(text)
-    cfg = _cfg()
-    obj, err = assemble(asm, cfg)
-    if err or obj is None:
-        if _DEBUG_SHOWN < _DEBUG_LIMIT:
-            with _LOCK:
-                if _DEBUG_SHOWN < _DEBUG_LIMIT:
-                    _DEBUG_SHOWN += 1
-                    print(f"[reward-debug #{_DEBUG_SHOWN}] vid={vid} asm_fail={err.message[:200] if err else 'None'!r}", flush=True)
-                    print(f"[reward-debug] raw[:300]={text[:300]!r}", flush=True)
-                    print(f"[reward-debug] extracted[:300]={asm[:300]!r}", flush=True)
-        cleanup_temp_dirs()
-        with _LOCK:
-            _CACHE[k] = e
-        return e
-    e.assembles = True
-    e.runs = run_correctness_qemu(obj, TestCase(inputs=(), expected=None), cfg)
-    if e.runs:
-        bc = _bcy(vid, basm)
-        try:
-            rep = run_mca(asm, cfg.mca_bin, cfg.mcpu)
-            e.speedup = bc / max(rep.total_cycles, 1)
-        except Exception:
-            e.speedup = 0.0
-    cleanup_temp_dirs()
-    with _LOCK:
-        _CACHE[k] = e
-    return e
-
-def _prep(completions, kwargs):
-    texts = [c[-1]["content"] if isinstance(c, list) else str(c) for c in (completions or [])]
-    n = len(texts)
-    vids = list(kwargs.get("variant_id") or [""] * n)
-    basms = list(kwargs.get("baseline_asm") or [""] * n)
-    if len(vids) == 1 and n > 1:
-        vids = vids * n
-        basms = basms * n
-    return texts, vids, basms
-
-def syntax_reward(prompts=None, completions=None, **kwargs) -> list[float]:
-    _ = prompts
-    texts, vids, basms = _prep(completions, kwargs)
-    return [1.0 if _run(t, v, b).assembles else 0.0 for t, v, b in zip(texts, vids, basms)]
-
-_FORMAT_DEBUG_STEP = 0
-_FORMAT_DEBUG_LIMIT = 6
-
-def format_reward(prompts=None, completions=None, **kwargs) -> list[float]:
-    global _FORMAT_DEBUG_STEP
-    _ = prompts
-    texts, _, _ = _prep(completions, kwargs)
-    if _FORMAT_DEBUG_STEP < _FORMAT_DEBUG_LIMIT and texts:
-        _FORMAT_DEBUG_STEP += 1
-        head = texts[0]
-        print(f"[completion-debug step={_FORMAT_DEBUG_STEP}] len={len(head)} first 500 chars:", flush=True)
-        print(repr(head[:500]), flush=True)
-        print(f"[completion-debug step={_FORMAT_DEBUG_STEP}] last 200 chars:", flush=True)
-        print(repr(head[-200:]), flush=True)
-    scores = []
-    for text in texts:
-        lowered = text.lower()
-        has_open = "<assembly>" in lowered
-        has_close = "</assembly>" in lowered
-        body = extract_assembly(text)
-        body_len = len(re.sub(r"\s", "", body))
-        has_prose = any(m in lowered for m in ("```", "<think>", "explain", "analysis"))
-        score = 0.0
-        if has_open:
-            score += 0.1
-        if has_close:
-            score += 0.1
-        if body_len >= 20:
-            score += 0.1
-        if has_prose:
-            score -= 0.05
-        scores.append(max(0.0, score))
-    return scores
-
-def correctness_reward(prompts=None, completions=None, **kwargs) -> list[float]:
-    _ = prompts
-    texts, vids, basms = _prep(completions, kwargs)
-    return [1.0 if _run(t, v, b).runs else 0.0 for t, v, b in zip(texts, vids, basms)]
-
-def speedup_reward(prompts=None, completions=None, **kwargs) -> list[float]:
-    _ = prompts
-    texts, vids, basms = _prep(completions, kwargs)
-    scores = []
-    for t, v, b in zip(texts, vids, basms):
-        e = _run(t, v, b)
-        scores.append(max(0.0, e.speedup - 1.0) if e.runs else 0.0)
-    return scores
-
-class LiveRewardFn:
-    @classmethod
-    def build(cls) -> "LiveRewardFn":
-        return cls()
-    def __call__(self, prompts=None, completions=None, **kwargs) -> list[float]:
-        texts, vids, basms = _prep(completions, kwargs)
-        out = []
-        for t, v, b in zip(texts, vids, basms):
-            e = _run(t, v, b)
-            if not e.assembles: out.append(0.0)
-            elif not e.runs: out.append(0.1)
-            else: out.append(max(0.0, e.speedup - 1.0))
-        return out
-REWARDPY
-echo "[hotpatch] reward_fn.py replaced (both-tag format_reward)"
+# 3-4. kaggle/dataset.py + kaggle/reward_fn.py: from clone (push before job; no heredoc)
+echo "[hotpatch] using kaggle/dataset.py + kaggle/reward_fn.py from clone"
 
 # Run toolchain setup
 [ "${HF_RUN_SETUP:-1}" = "1" ] && [ -f kaggle/setup.sh ] && bash kaggle/setup.sh

kaggle/reward_fn.py

@@ -1,9 +1,13 @@
 """Reward functions for GRPOTrainer - 3 independent callables for GDPO.
 
-syntax_reward     : +1.0 if assembly parses, 0.0 otherwise
+syntax_reward     : +1.0 if assembly parses, **-1.0** otherwise (was 0.0)
 correctness_reward: +1.0 if QEMU executes without crash, 0.0 otherwise
 speedup_reward    : (baseline_cycles / agent_cycles) - 1.0; 0.0 if !runs
 
+Stronger **invalid** syntax signal (negative) avoids PPO/GRPO treating “broken” and
+“almost right” as the same plateau. Pre-``as`` fast-reject (unclosed block comments,
+obvious high-level code slop) saves subprocess budget and fails before GNU as.
+
 Shared _Entry cache avoids 3x verifier calls per completion within a batch.
 Cache is module-level and persistent (verifier is deterministic, results stable).
 GDPO: pass [syntax_reward, correctness_reward, speedup_reward] to reward_funcs
@@ -26,6 +30,20 @@ from arm_gym.verifier import VerifierConfig, assemble, cleanup_temp_dirs
 _ASM_RE = re.compile(r"<assembly>(.*?)</assembly>", re.DOTALL | re.IGNORECASE)
 _THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
 
+# Syntax gate: separate “dead” from “valid” in RL (not flat zeros).
+SYNTAX_REWARD_OK: float = 1.0
+SYNTAX_REWARD_FAIL: float = -1.0
+
+# Substrings that almost never appear in hand-written A64 gas but show up in C/Java/JS drift.
+_CODE_SLOP = re.compile(
+    r"(?i)\b("
+    r"import|package|#include|#pragma|"
+    r"class\s|def\s|void\s+main|public\s+static|"
+    r"console\.|React\.|"
+    r"System\.|using\s+namespace"
+    r")\b",
+)
+
 
 def extract_assembly(text: str) -> str:
     """Extract assembly body from text.
@@ -58,6 +76,43 @@ def _clean_assembly(text: str) -> str:
     return text.strip()
 
 
+def preprocess_assembly(asm: str) -> tuple[str, str | None]:
+    """Sanitize for GNU ``as`` and **fast-reject** obvious invalid text before ``assemble()``.
+
+    * Unclosed C block comments → reject (GAS: "end of file in multiline comment").
+    * Obvious high-level / JS / C drift → reject (saves `as` + clearer RL signal than junk errors).
+    * ``//`` → ``@`` line comments (A64 gas line comment; plain ``as`` does not treat ``//`` like gcc).
+    * Smart quotes / mojibake quotes → ASCII (assembler is strict).
+    * Trailing newline (silences warning; some gas versions are picky).
+
+    We **do not** use a naive ``[a-zA-Z]{6,}`` filter — valid mnemonics like
+    ``umaddl``, ``cfinv``, ``st1`` would false-positive.
+
+    Returns ``(text, None)`` if OK, or ``("", reason)`` to skip ``assemble()``.
+    """
+    if not (asm and asm.strip()):
+        return "", "empty_body"
+
+    s = (
+        asm.replace("\u201c", '"')
+        .replace("\u201d", '"')
+        .replace("\u2018", "'")
+        .replace("\u2019", "'")
+    )
+    # C++-style line comments: map to GAS line comment; avoids "junk" on `/` for plain `.s`.
+    s = s.replace("//", "@")
+
+    if s.count("/*") != s.count("*/"):
+        return "", "unclosed_block_comment"
+
+    if _CODE_SLOP.search(s):
+        return "", "lexical:high_level_code"
+
+    if not s.endswith("\n"):
+        s += "\n"
+    return s, None
+
+
 @dataclass
 class _Entry:
     assembles: bool = False
@@ -114,10 +169,30 @@ def _run(text: str, vid: str, basm: str) -> _Entry:
     # Compute outside lock - subprocess calls can be slow.
     # Two threads racing on the same key both compute and write; result is identical.
     e = _Entry()
-    asm = extract_assembly(text)
+    raw_extracted = extract_assembly(text)
+    asm, pre_r = preprocess_assembly(raw_extracted)
     cfg = _cfg()
 
-    # Gate 1: assemble
+    # Gate 1: assemble (skip subprocess if fast-reject already failed)
+    if pre_r is not None:
+        if _DEBUG_SHOWN < _DEBUG_LIMIT:
+            with _LOCK:
+                if _DEBUG_SHOWN < _DEBUG_LIMIT:
+                    _DEBUG_SHOWN += 1
+                    print(
+                        f"[reward-debug #{_DEBUG_SHOWN}] vid={vid} "
+                        f"pre_reject={pre_r}",
+                        flush=True,
+                    )
+                    print(f"[reward-debug] raw[:300]={text[:300]!r}", flush=True)
+                    print(
+                        f"[reward-debug] extracted[:300]={raw_extracted[:300]!r}",
+                        flush=True,
+                    )
+        with _LOCK:
+            _CACHE[k] = e
+        return e
+
     obj, err = assemble(asm, cfg)
     if err or obj is None:
         if _DEBUG_SHOWN < _DEBUG_LIMIT:
@@ -130,7 +205,10 @@ def _run(text: str, vid: str, basm: str) -> _Entry:
                         flush=True,
                     )
                     print(f"[reward-debug] raw[:300]={text[:300]!r}", flush=True)
-                    print(f"[reward-debug] extracted[:300]={asm[:300]!r}", flush=True)
+                    print(
+                        f"[reward-debug] extracted[:300]={raw_extracted[:300]!r}",
+                        flush=True,
+                    )
         cleanup_temp_dirs()
         with _LOCK:
             _CACHE[k] = e
@@ -140,7 +218,7 @@ def _run(text: str, vid: str, basm: str) -> _Entry:
     # Gate 2: QEMU run-without-crash
     e.runs = run_correctness_qemu(obj, TestCase(inputs=(), expected=None), cfg)
 
-    # Gate 3: MCA speedup
+    # Gate 3: MCA speedup (use sanitized asm, same as ``assemble``)
     if e.runs:
         bc = _bcy(vid, basm)
         try:
@@ -167,38 +245,47 @@ def _prep(completions, kwargs):
 
 
 def syntax_reward(prompts=None, completions=None, **kwargs) -> list[float]:
-    """Gate 1: +1.0 if assembly parses, 0.0 otherwise."""
+    """Gate 1: ``SYNTAX_REWARD_OK`` if assembly parses, ``SYNTAX_REWARD_FAIL`` otherwise."""
     _ = prompts
     texts, vids, basms = _prep(completions, kwargs)
-    return [1.0 if _run(t, v, b).assembles else 0.0 for t, v, b in zip(texts, vids, basms)]
+    return [
+        (SYNTAX_REWARD_OK if _run(t, v, b).assembles else SYNTAX_REWARD_FAIL)
+        for t, v, b in zip(texts, vids, basms)
+    ]
+
+
+FORMAT_REWARD_MISSING_CLOSE: float = -1.0
 
 
 def format_reward(prompts=None, completions=None, **kwargs) -> list[float]:
     """Shaping reward for well-formed completions.
 
-    The model should emit both ``<assembly>`` and ``</assembly>`` tags.
-      +0.1  has ``<assembly>`` open tag
-      +0.1  has ``</assembly>`` close tag
-      +0.1  has non-trivial assembly body (>= 20 non-whitespace chars between tags)
-      -0.05 contains prose markers (code fences, think tags, explanations)
-    Max score: 0.3 for a perfectly formed completion.
+    If ``</assembly>`` is **missing** → ``FORMAT_REWARD_MISSING_CLOSE`` (RL signal
+    to separate unclosed from weak-but-closed). Otherwise small positive parts.
+
+    Prompt may prefill ``<assembly>`` so the assistant body often omits the open
+    tag; we still require a closing tag and use ``extract_assembly`` for body.
     """
     _ = prompts
     texts, _, _ = _prep(completions, kwargs)
     scores = []
     for text in texts:
         lowered = text.lower()
+        if "</assembly>" not in lowered:
+            scores.append(FORMAT_REWARD_MISSING_CLOSE)
+            continue
+
         has_open = "<assembly>" in lowered
-        has_close = "</assembly>" in lowered
         body = extract_assembly(text)
         body_len = len(re.sub(r"\s", "", body))
+        # Prefill: open tag only in prompt — credit body as soon as it exists.
+        has_open_effective = has_open or body_len >= 1
         has_prose = any(marker in lowered
                         for marker in ("```", "<think>", "explain", "analysis"))
         score = 0.0
-        if has_open:
-            score += 0.1
-        if has_close:
+        if has_open_effective:
             score += 0.1
+        score += 0.1  # has_close: we only get here if </assembly> present
         if body_len >= 20:
             score += 0.1
         if has_prose:
@@ -242,7 +329,7 @@ class LiveRewardFn:
         for t, v, b in zip(texts, vids, basms):
             e = _run(t, v, b)
             if not e.assembles:
-                out.append(0.0)
+                out.append(SYNTAX_REWARD_FAIL)
             elif not e.runs:
                 out.append(0.1)
             else:

scripts/sync_hf_space.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Force-push current branch to Hugging Face Space main, rewriting plot PNGs to Git LFS
+# first so the Hub accepts the pack (raw PNG blobs are rejected).
+#
+#   git fetch origin && git checkout main && git pull
+#   export HF_TOKEN=hf_...   # or use git credential for huggingface.co
+#   ./scripts/sync_hf_space.sh
+#
+set -euo pipefail
+ROOT="$(git rev-parse --show-toplevel)"
+cd "$ROOT"
+
+HF_USER="${HF_USER:-kaori02}"
+HF_SPACE="${HF_SPACE:-arm-gym}"
+if [[ -n "${HF_TOKEN:-}" ]]; then
+  REMOTE_URL="https://${HF_USER}:${HF_TOKEN}@huggingface.co/spaces/${HF_USER}/${HF_SPACE}"
+else
+  REMOTE_URL="https://huggingface.co/spaces/${HF_USER}/${HF_SPACE}"
+fi
+
+REF=$(git symbolic-ref -q HEAD) || { echo "error: need a branch (detached HEAD not supported)" >&2; exit 1; }
+
+git lfs install
+git lfs migrate import --include="colab/results/plots/*.png" --include-ref="$REF"
+
+echo "[sync_hf_space] pushing HEAD -> ${HF_USER}/${HF_SPACE}:main (force)"
+git push --force "$REMOTE_URL" HEAD:main

tests/test_preprocess_assembly.py

@@ -0,0 +1,53 @@
+"""Unit tests for ``preprocess_assembly`` (no toolchain)."""
+
+from __future__ import annotations
+
+import pytest
+
+from kaggle.reward_fn import SYNTAX_REWARD_FAIL, SYNTAX_REWARD_OK, preprocess_assembly, syntax_reward
+
+
+def test_preprocess_appends_newline() -> None:
+    s, err = preprocess_assembly("mov x0, #0")
+    assert err is None
+    assert s.endswith("\n")
+
+
+def test_preprocess_maps_double_slash_to_at() -> None:
+    s, err = preprocess_assembly("mov x0, x1 // foo\n")
+    assert err is None
+    assert "//" not in s
+    assert "@" in s
+
+
+def test_preprocess_rejects_unclosed_block_comment() -> None:
+    _, err = preprocess_assembly("mov x0, x1 /* start\n")
+    assert err == "unclosed_block_comment"
+
+
+def test_preprocess_rejects_code_slop_import() -> None:
+    _, err = preprocess_assembly("import java.util\nmov x0, x0\n")
+    assert err == "lexical:high_level_code"
+
+
+def test_preprocess_allows_valid_mnemonic_umaddl() -> None:
+    # Do not use naive "long English word" heuristics — umaddl is 6+ letters.
+    s, err = preprocess_assembly("umaddl x0, w1, w2, w3\n")
+    assert err is None
+    assert "umaddl" in s
+
+
+def test_syntax_reward_negative_on_garbage(monkeypatch: pytest.MonkeyPatch) -> None:
+    from kaggle import reward_fn as rf
+
+    monkeypatch.setattr(rf, "_CACHE", {})
+    monkeypatch.setattr(rf, "_DEBUG_SHOWN", 0)
+    # No valid assembly, pre-reject or assemble fail
+    bad = "<assembly>import numpy as np\n</assembly>"
+    r = syntax_reward(completions=[[{"content": bad}]], variant_id=["t"], baseline_asm=[""])
+    assert r == [SYNTAX_REWARD_FAIL]
+
+
+def test_syntax_reward_ok_constant() -> None:
+    assert SYNTAX_REWARD_OK == 1.0
+    assert SYNTAX_REWARD_FAIL == -1.0

tests/test_training_prompt_contract.py

@@ -36,11 +36,14 @@ def test_format_reward_prefers_closed_assembly_blocks() -> None:
         ]
     )
 
-    assert rewards == [0.2, 0.0]
+    assert rewards == [0.2, -1.0]
 
 
 def test_user_prompt_includes_baseline_fallback() -> None:
     prompt = user_prompt("int f(void){return 0;}", "f:\n\tret\n")
 
-    assert "copy the baseline assembly exactly" in prompt
-    assert "...AArch64 assembly only..." not in prompt
+    assert "Generate ONLY AArch64" in prompt
+    assert "C code:" in prompt
+    assert "Baseline assembly" in prompt
+    assert prompt.rstrip().endswith("<assembly>")
+    assert "ret" in prompt