Refactor mxmoe

README.md

@@ -368,7 +368,6 @@ See `examples/` for more examples including tiled copy (`02-tiledCopy.py`), tile
 | **Blockscale GEMM** | `test_blockscale_preshuffle_gemm.py` | Blockscale preshuffle GEMM |
 | **HGEMM Split-K** | `test_hgemm_splitk.py` | FP16 GEMM split-K |
 | **MoE GEMM** | `test_moe_gemm.py` | MoE 2-stage (gate/up + reduce) |
-| **MoE Blockscale** | `test_moe_blockscale.py` | MoE blockscale 2-stage |
 | **MoE Reduce** | `test_moe_reduce.py` | MoE reduce kernel |
 | **PagedAttention** | `test_pa.py` | Paged attention decode (FP8) — *WIP perf tuning* |
 | **FlashAttention** | `test_flash_attn_fwd.py` | Flash attention — *WIP perf tuning* |

docs/api/kernels.rst

@@ -21,8 +21,7 @@ MoE (Mixture-of-Experts) Kernels
 - ``kernels.moe.moe_gemm_2stage`` -- MoE GEMM with 2-stage pipeline (stage1 + stage2). Also
   provides the MoE reduction (sum over the topk dimension, ``Y[t, d] = sum(X[t, :, d])``),
   compiled via ``compile_moe_reduction()``.
-- ``kernels.moe.mixed_moe_gemm_2stage`` -- Mixed-precision MoE GEMM
-- ``kernels.moe.moe_blockscale_2stage`` -- MoE with block-scale quantization (MXFP4)
+- ``kernels.moe.mxfp_moe`` -- Fused a4w4 / a8w4 MoE 2-stage GEMM (device-side fp4 re-quant)
 
 Paged Attention
 ----------------

docs/architecture_guide.md

@@ -103,8 +103,7 @@ FlyDSL/
 │   │   └── fused_rope_cache_kernel.py # Fused RoPE + KV cache
 │   ├── moe/                          # Mixture-of-experts kernels
 │   │   ├── moe_gemm_2stage.py        # MoE GEMM (2-stage gate/up + reduce)
-│   │   ├── moe_blockscale_2stage.py  # MoE Blockscale GEMM
-│   │   ├── mixed_moe_gemm_2stage.py  # Mixed-precision MoE GEMM
+│   │   ├── mxfp_moe/               # Fused a4w4/a8w4 MoE 2-stage (device fp4 re-quant)
 │   │   ├── moe_sorting_kernel.py     # MoE token sorting
 │   │   └── topk_gating_softmax_kernel.py # Top-k gating softmax
 │   ├── mma/                          # Shared MMA pipeline helpers

docs/prebuilt_kernels_guide.md

@@ -348,7 +348,6 @@ What operation do you need?
 │
 ├── MoE (Mixture of Experts)
 │   ├── Blockscale MoE (gate+up+reduce)
-│   │   └── → kernels/moe/moe_blockscale_2stage.py
 │   └── Standard MoE (fp8/f16/bf16/int8/int4)
 │       └── → kernels/moe/moe_gemm_2stage.py
 │
@@ -368,8 +367,7 @@ What operation do you need?
 | `kernels/gemm/blockscale_preshuffle_gemm.py` | Blockscale GEMM |
 | `kernels/gemm/hgemm_splitk.py` | FP16 GEMM split-K |
 | `kernels/moe/moe_gemm_2stage.py` | MoE GEMM 2-stage (gate/up + reduce) |
-| `kernels/moe/moe_blockscale_2stage.py` | MoE Blockscale 2-stage |
-| `kernels/moe/mixed_moe_gemm_2stage.py` | Mixed-precision MoE GEMM |
+| `kernels/moe/mxfp_moe/` | Fused a4w4/a8w4 MoE 2-stage GEMM (device fp4 re-quant) |
 | `kernels/attention/pa_decode_fp8.py` | Paged attention decode (FP8) |
 | `kernels/attention/flash_attn_generic.py` | FlashAttention generic fallback |
 | `kernels/attention/flash_attn_gfx950.py` | FlashAttention gfx950 bf16/f16 fast path |
@@ -398,7 +396,6 @@ What operation do you need?
 | `tests/kernels/test_blockscale_preshuffle_gemm.py` | Blockscale GEMM |
 | `tests/kernels/test_hgemm_splitk.py` | FP16 GEMM split-K |
 | `tests/kernels/test_moe_gemm.py` | MoE GEMM |
-| `tests/kernels/test_moe_blockscale.py` | MoE Blockscale GEMM |
 | `tests/kernels/test_moe_reduce.py` | MoE reduce kernel |
 | `tests/kernels/test_pa.py` | Paged attention decode |
 | `tests/kernels/test_flash_attn_fwd.py` | FlashAttention |

docs/testing_benchmarking_guide.md

@@ -69,7 +69,6 @@ Full end-to-end tests: compile FlyDSL kernels, execute on GPU, validate against
 | `test_preshuffle_gemm.py` | GEMM | Preshuffle MFMA GEMM (fp8/int8/fp16/bf16) |
 | `test_blockscale_preshuffle_gemm.py` | GEMM | Block-scale (MXFP4) preshuffle GEMM |
 | `test_moe_gemm.py` | MoE GEMM | Mixture-of-Experts GEMM |
-| `test_moe_blockscale.py` | MoE | MoE with block-scale quantization |
 | `test_moe_reduce.py` | MoE Reduce | MoE reduction kernel |
 | `test_pa.py` | Paged Attn | Paged attention decode |
 | `test_quant.py` | Quantization | Quantization ops |

kernels/common/tensor_shim.py

@@ -19,13 +19,23 @@
 def _run_compiled(exe, *args):
     """First call: ``flyc.compile(exe, *args)`` compiles **and** executes the kernel.
     Subsequent calls: fast dispatch via the cached ``CompiledFunction``.
+
+    A failed cold compile can leave an MLIR ``Context`` open on the stack; unwind
+    it before re-raising so the next attempt starts clean.
     """
     cf = getattr(exe, "_cf", None)
-    if cf is None:
-        cf = flyc.compile(exe, *args)
-        exe._cf = cf
-    else:
+    if cf is not None:
         cf(*args)
+        return
+    try:
+        exe._cf = flyc.compile(exe, *args)
+    except Exception:
+        try:
+            while ir.Context.current is not None:
+                ir.Context.current.__exit__(None, None, None)
+        except Exception:
+            pass
+        raise
 
 
 def _to_raw(v):