"""Model-free correctness receipts from the pinned MTPLX custom kernels.

No timings, generation, weights download, or replacement implementations.
Run with the pinned runtime and PYTHONPATH pointing at the MTPLX checkout.
"""

import hashlib
import json
import os
from pathlib import Path
import runpy
import sys

import mlx.core as mx
from mlx_lm.models import gated_delta
import numpy as np
from mtplx.kernels.gdn_conv_norm import fused_gdn_conv_norm, fused_gdn_conv_norm_rows
from mtplx.kernels.gdn_step_fused import fused_gdn_step
from mtplx.kernels.gdn_out_fused import fused_gdn_out
from mtplx.kernels.hyper_connection import fused_hyper_read
from mtplx.kernels.hyper_connection_v3 import fused_hyper_read_v3
from mtplx.kernels import hyper_connection


def pattern(count, salt):
    return mx.array((((np.arange(count, dtype=np.int64) * 17 + salt * 13) % 257)
                     - 128).astype(np.float32) / 128, dtype=mx.bfloat16)


def qsa_attention_module_fixtures(*, profile_filter=None):
    """Whole unchanged Attention, ongoing KV/indexer state and actual consumers."""
    from contextlib import nullcontext
    from mtplx.models import qwen4_exp as model
    from mtplx.attention_context import attention_phase, vision_rope
    from mtplx.kernels import qsa_flash_skip, qsa_prefill_flash
    assert mx.__version__ == "0.32.2"
    defaults=dict(fused=True,compiled=True,prefill=True,phase_prefill=True,
        prefill_min_rows=32,prefill_compile_rows=2048,prefill_min_context=2049,
        prefill_workspace=128*1024*1024,tile_rows=0,flash=True,gather_decode=False,
        gather=True,gather_max_rows=8,gather_min_context=0,tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0")
    envs=dict(fused="MTPLX_FUSED_QSA_INDEXER",compiled="MTPLX_COMPILED_QSA_INDEXER",prefill="MTPLX_QSA_PREFILL",
        prefill_min_rows="MTPLX_QSA_PREFILL_MIN_ROWS",prefill_compile_rows="MTPLX_QSA_PREFILL_COMPILE_ROWS",
        prefill_min_context="MTPLX_QSA_PREFILL_MIN_CONTEXT",flash="MTPLX_QSA_FLASH",gather_decode="MTPLX_QSA_GATHER_DECODE",
        gather="MTPLX_QSA_GATHER",gather_max_rows="MTPLX_QSA_GATHER_MAX_ROWS",gather_min_context="MTPLX_QSA_GATHER_MIN_CONTEXT",tile_rows="MTPLX_QSA_SCORE_TILE_ROWS")
    saved=os.environ.copy()
    profiles=[(0,(1,4,33),4,{},"none",False,False,2049,True),
        (0,(2048,1),3,{},"none",False,False,2049,True),
        (2051,(1,4,33),4,{},"none",False,False,2049,True),
        (2051,(33,1,4),0,dict(compiled=False,flash=False,gather=False),"none",False,False,32768,True),
        (2051,(33,1,4),3,{},"none",False,False,32768,False),
        (2051,(1,4),0,dict(compiled=False,fused=False,flash=False,gather_decode=True),"none",False,False,32768,True),
        (32768,(2048,1,4),4,{},"none",False,False,32768,True),
        (2051,(1,4,2,1),4,{},"table",False,False,2049,True),
        (0,(3,1,4),0,{},"none",True,True,2049,True),
        (2051,(1,4),3,{},"no_axes",False,False,2049,True)]
    trace=[];hooks=[]
    def hook(owner,name,label):
        original=getattr(owner,name)
        def call(*args,**kwargs):
            trace.append(label);return original(*args,**kwargs)
        hooks.append((owner,name,original));setattr(owner,name,call)
    hook(qsa_flash_skip,"qsa_flash_skip","flash")
    hook(qsa_prefill_flash,"qsa_prefill_flash","prefill_flash")
    hook(model,"_qsa_rows_gather_attention","rows_gather")
    hook(model,"_qsa_prefill_gather_attention","prefill_gather")
    hook(model,"_qsa_blocks_to_dense_mask","dense_mask")
    hook(mx.fast,"scaled_dot_product_attention","sdpa")
    try:
        for profile,(prefix,steps,fused,override,vision,no_indexer,biased,flash_min,gather) in enumerate(profiles):
            if profile_filter is not None and profile!=profile_filter:continue
            opt=defaults|override
            os.environ.update({env:str(int(opt[key])) for key,env in envs.items()})
            os.environ.update(MTPLX_QSA_PREFILL_SCORE_MB="128",MTPLX_QSA_PREFILL_FLASH_MIN_CONTEXT=str(flash_min),MTPLX_QSA_PREFILL_GATHER=str(int(gather)),MTPLX_QSA_PREFILL_GATHER_TILE="64")
            args=model.TextArgs(hidden_size=64,indexer_n_heads=0 if no_indexer else 4,
                mrope_section=[8,12,12] if vision=="table" else None,mrope_interleaved=True)
            attn=model.Attention(args)
            for name,n,k,salt in [("q_proj",12288,64,400),("k_proj",512,64,402),("v_proj",512,64,404),("o_proj",64,6144,406)]:
                layer=quantized_linear_fixture(n,k,4,64,salt)
                if biased:layer.bias=pattern(n,salt+40)/64
                setattr(attn,name,layer)
            attn.q_norm.weight=pattern(256,410);attn.k_norm.weight=pattern(256,411)
            if attn.indexer is not None:
                attn.indexer.index_qk_proj=quantized_linear_fixture(640,64,4,64,408)
                attn.indexer.q_layernorm.weight=pattern(128,72);attn.indexer.k_layernorm.weight=pattern(128,73)
            if fused:
                parts=[attn.q_proj,attn.k_proj,attn.v_proj]
                if fused==4:parts.append(attn.indexer.index_qk_proj)
                splits=np.cumsum([a.weight.shape[0] for a in parts[:-1]]).tolist()
                attn.qkv_fused=model._FusedGDNInProj(*(mx.concatenate([getattr(a,name) for a in parts],axis=0) for name in ("weight","scales","biases")),64,4,"affine",splits)
                for name in ("q_proj","k_proj","v_proj"):attn.pop(name,None)
                if fused==4:attn.indexer.pop("index_qk_proj",None)
            cache=model.QSACache()
            if prefix:
                cache.write_raw(pattern(prefix*128,88).reshape(1,prefix,128))
                attn.indexer._extend_pooled(cache,prefix)
                kv=pattern(prefix*512,89).reshape(1,prefix,2,256).transpose(0,2,1,3)
                cache.kv.update_and_fetch(kv,kv)
                mx.eval(*cache.state)
            table=mx.array(np.stack([np.arange(prefix+5)//3,np.arange(prefix+5)%17,np.arange(prefix+5)%11]),mx.int32) if vision=="table" else None
            for step,rows in enumerate(steps):
                if step==2 and profile in (2,7):assert cache.trim(2)==2
                pos=cache.offset
                trace.clear();before=model.qsa_prefill_engagement()
                with attention_phase("prefill"), vision_rope(table,-2) if vision!="none" else nullcontext():
                    result=attn(pattern(rows*64,420+step).reshape(1,rows,64),cache)
                counts={k:v-before.get(k,0) for k,v in model.qsa_prefill_engagement().items() if v!=before.get(k,0)}
                state=[cache.raw_keys,cache.pooled,cache.pooled_f32_t,cache.kv.keys,cache.kv.values]
                leaves=[result,*[a for a in state if a is not None]]
                core=None if attn.indexer is None else attn.indexer._compiled_indexer_core
                emit(f"qsa_attention_p{profile}_s{step}",leaves,profile=profile,step=step,rows=rows,pos=pos,prefix=prefix,
                    fused=fused,vision=vision,no_indexer=no_indexer,biased=biased,flash_min=flash_min,prefill_gather=gather,
                    options=opt,consumers=trace.copy(),counts=counts,offset=cache.offset,pooled_len=cache.pooled_len,
                    state_present=[a is not None for a in state],shapes=[a.shape for a in leaves],
                    core=None if core is None else core.to_dict(),frequencies=attn._inv_freq.tolist())
    finally:
        for owner,name,original in reversed(hooks):setattr(owner,name,original)
        for key in set(envs.values())|{"MTPLX_QSA_PREFILL_SCORE_MB","MTPLX_QSA_PREFILL_FLASH_MIN_CONTEXT","MTPLX_QSA_PREFILL_GATHER","MTPLX_QSA_PREFILL_GATHER_TILE"}:
            if key in saved:os.environ[key]=saved[key]
            else:os.environ.pop(key,None)


def qsa_indexer_commit_fixtures():
    """Unmodified indexer core lifecycle, ongoing cache commit and return API."""
    from mtplx.models.qwen4_exp import QSAIndexer, TextArgs, QSACache
    import mlx.nn as nn
    assert mx.__version__ == "0.32.2"
    saved = os.environ.get("MTPLX_QSA_PREFILL_SCORE_MB")
    steps = [(1,"update_only"),(3,"update_only"),(33,"prefill_blocks"),
        (1,"blocks"),(4,"row_tokens"),(4,"dense_mask"),(221,"prefill_blocks"),
        (1,"blocks"),(4,"row_tokens"),(4,"dense_mask"),(4,"row_tokens"),
        (1,"blocks"),(4,"dense_mask")]
    try:
        for quantized in (False, True):
            os.environ["MTPLX_QSA_PREFILL_SCORE_MB"] = "128"
            idx=QSAIndexer(TextArgs());idx.block_topk=8
            idx.q_layernorm.weight=pattern(128,72);idx.k_layernorm.weight=pattern(128,73)
            if quantized:idx.index_qk_proj=quantized_linear_fixture(640,128,4,64,75)
            else:
                idx.index_qk_proj=nn.Linear(128,640,bias=False)
                idx.index_qk_proj.weight=pattern(640*128,231).reshape(640,128)
            cache=QSACache()
            for step,(rows,mode) in enumerate(steps):
                # Same bytes, new Python array identity: every leaf matters.
                if step==2:
                    idx._fused_score_scratch_bytes=1024*1024
                    os.environ["MTPLX_QSA_PREFILL_SCORE_MB"]="2"
                if step==3:idx.q_layernorm.weight=mx.array(idx.q_layernorm.weight)
                if step==4:idx.k_layernorm.weight=mx.array(idx.k_layernorm.weight)
                if step==5:idx._inv_freq=mx.array(idx._inv_freq)
                if step==6:idx.index_qk_proj.weight=mx.array(idx.index_qk_proj.weight)
                if step==7 and quantized:idx.index_qk_proj.scales=mx.array(idx.index_qk_proj.scales)
                if step==8 and quantized:idx.index_qk_proj.biases=mx.array(idx.index_qk_proj.biases)
                if step==9:cache.reserve_indexer_capacity(raw_capacity=2048,pooled_capacity=512)
                if step==10:assert cache.trim(3)==3
                pos=cache.offset
                if step==11:
                    old_core=idx._compiled_indexer_core;old_seal=idx._compiled_indexer_parameter_signature
                    valid=idx.q_layernorm.weight
                    idx.q_layernorm.weight=mx.zeros((127,),mx.bfloat16)
                    try:idx._get_compiled_indexer_core()
                    except ValueError:pass
                    else:raise AssertionError("invalid replacement must fail")
                    assert idx._compiled_indexer_core is old_core and idx._compiled_indexer_parameter_signature==old_seal
                    idx.q_layernorm.weight=valid
                hidden=pattern(rows*128,300+step).reshape(1,rows,128)
                supplied=step%2==0
                qk=pattern(rows*704,320+step).reshape(1,rows,704)[...,:640] if supplied else None
                if cache.pooled is not None:cache.pooled_f32_view(cache.pooled_len)
                old_raw,old_pool=cache.raw_keys,cache.pooled
                old_arrays=[a for a in (old_raw,old_pool) if a is not None]
                mx.eval(*old_arrays)
                result=idx._call_rows_compiled(hidden,pos,cache,qk,mode=mode)
                assert cache.offset==pos and cache.pooled_f32_t is None
                assert cache.raw_keys is not old_raw and cache.pooled is not old_pool
                if result is None:kind,values,tail="none",[],None
                elif isinstance(result,tuple):
                    kind=result[0];tail=result[2] if kind=="flash" else None
                    values=[result[1]] if kind=="flash" else list(result[1:])
                else:kind,values,tail="dense",[result],None
                outputs=[*values,cache.raw_keys,cache.pooled,*old_arrays]
                emit(f"qsa_indexer_commit_q{int(quantized)}_s{step}",outputs,
                    quantized=quantized,step=step,rows=rows,pos=pos,mode=mode,supplied=supplied,
                    kind=kind,tail=tail,shapes=[a.shape for a in outputs],
                    pooled_len=cache.pooled_len,report=idx._compiled_indexer_core.to_dict(),
                    frequencies=idx._inv_freq.tolist())
                cache.kv.offset=pos+rows  # Attention owns this later frontier.
    finally:
        if saved is None:os.environ.pop("MTPLX_QSA_PREFILL_SCORE_MB",None)
        else:os.environ["MTPLX_QSA_PREFILL_SCORE_MB"]=saved


def qsa_core_host_fixtures():
    """Original Core validation, engagement and replay; real compile/eval."""
    from mtplx.kernels.qsa_indexer_compile import QSACompiledIndexerCore
    from mtplx.models.qwen4_exp import QSAIndexer, TextArgs
    assert mx.__version__ == "0.32.2"
    norm=pattern(128,72);freq=QSAIndexer(TextArgs())._inv_freq
    weight=pattern(640*128,231).reshape(640,128)
    def make(project=lambda x:x@weight.T):
        return QSACompiledIndexerCore(n_heads=4,kv_heads=1,head_dim=128,block_topk=8,
            compress_ratio=4,q_norm_weight=norm,k_norm_weight=norm,inv_freq=freq,rms_norm_eps=1e-6,
            project_qk=project,minimum_raw_capacity=16,minimum_pooled_capacity=16,
            selector_scratch_bytes=1024,prefill_score_workspace_bytes=4096)
    core=make()
    cases=[(1,64,128,32,"blocks",False,None), (1,65,128,32,"blocks",False,None),
        (4,64,128,32,"row_tokens",False,None),(4,65,128,32,"row_tokens",False,None),
        (17,64,128,32,"prefill_blocks",False,None),(17,65,256,64,"prefill_blocks",False,None),
        (4,65,256,64,"dense_mask",True,None),(4,66,256,64,"update_only",True,None),
        (4,65,256,64,"dense_mask",False,300),(4,65,256,64,"dense_mask",False,None),
        (1,64,128,32,"blocks",False,None)]
    for i,(rows,pos,rc,pc,mode,hidden,dense) in enumerate(cases):
        width=128 if hidden else 640
        stride=width+(64 if i%2 else 0)
        source=pattern(rows*stride,232+i).reshape(1,rows,stride)[...,:width]
        raw=pattern(rc*128,250+i).reshape(1,rc,128)
        pool=pattern(pc*128,270+i).reshape(1,pc,128)
        front=[pos,pos+rows,(pos+rows)//4,pos//4]
        states=[mx.array(n,mx.int32) if i%2 else n for n in front]
        out=(core.select_hidden if hidden else core.select_qk_rows)(source,raw,pool,pos_start=states[0],total_tokens=states[1],logical_blocks=states[2],pooled_len=states[3],mode=mode,dense_output_tokens=dense)
        selected=[] if out.selection is None else list(out.selection) if isinstance(out.selection,tuple) else [out.selection]
        report=core.to_dict();report.pop("compiled_keys")
        emit(f"qsa_core_host_{i}",[*selected,out.raw_keys,out.pooled,out.pooled_len,out.offset],case=i,rows=rows,pos=pos,raw_capacity=rc,pooled_capacity=pc,mode=mode,hidden=hidden,dense=dense,frequencies=freq.tolist(),report=report)
    print(json.dumps(dict(kernel="qsa_core_report",report=core.to_dict()),separators=(",",":")),flush=True)
    invalid=[dict(pos=-1),dict(pos=127),dict(total=69),dict(total=-1),dict(logical=33),dict(logical=15),dict(frontier=0),dict(frontier=33),dict(raw_capacity=96),dict(raw_capacity=2),dict(pooled_capacity=0),dict(dense=100,mode="dense_mask"),dict(dense=1),dict(mode="prefill_blocks",rows=1),dict(width=639),dict(index_dtype="float32"),dict(index_dtype="int32",pos=127)]
    for i,change in enumerate(invalid):
        params=dict(pos=64,total=68,logical=17,frontier=16,raw_capacity=128,pooled_capacity=32,dense=None,mode="blocks",rows=4,width=640,index_dtype=None)|change
        x=mx.zeros((1,params["rows"],params["width"]),mx.bfloat16)
        raw=mx.zeros((1,params["raw_capacity"],128),mx.bfloat16);pool=mx.zeros((1,params["pooled_capacity"],128),mx.bfloat16)
        pos=mx.array(params["pos"],getattr(mx,params["index_dtype"])) if params["index_dtype"] else params["pos"]
        try:
            core.select_qk_rows(x,raw,pool,pos_start=pos,total_tokens=params["total"],logical_blocks=params["logical"],pooled_len=params["frontier"],mode=params["mode"],dense_output_tokens=params["dense"])
        except (ValueError,TypeError):ok=True
        else:ok=False
        print(json.dumps(dict(kernel=f"qsa_core_invalid_{i}",params=params,rejected=ok),separators=(",",":")),flush=True)
    bad=make(lambda x:x@weight[:639].T)
    for i in range(2):
        try:bad.select_hidden(mx.zeros((1,4,128),mx.bfloat16),mx.zeros((1,128,128),mx.bfloat16),mx.zeros((1,32,128),mx.bfloat16),pos_start=64,total_tokens=68,logical_blocks=17,pooled_len=16,mode="blocks")
        except ValueError:ok=True
        else:ok=False
        print(json.dumps(dict(kernel=f"qsa_core_failed_trace_{i}",rejected=ok,report=bad.to_dict()),separators=(",",":")),flush=True)
    print(json.dumps(dict(kernel="qsa_core_report_after_invalid",report=core.to_dict()),separators=(",",":")),flush=True)


def dynamic_slice_graph_fixtures(dtype_tag):
    """Original runtime dynamic slice/update; indices stay device arrays."""
    assert mx.__version__=="0.32.2"
    dtype=getattr(mx,dtype_tag)
    profiles=[((1,16,8),(1,3,8),(1,),(5,),"plain"),
        ((1,16,8),(1,3,8),(1,),(12,),"stride"),
        ((1,16,8),(1,3,8),(1,),(6,),"transpose"),
        ((1,16,8),(1,3,8),(1,),(5,),"broadcast"),
        ((3,7,8),(2,3,4),(0,1,2),(1,2,3),"plain"),
        ((2,3,4,5),(1,2,3,4),(0,1,2,3),(1,1,1,1),"stride"),
        ((1,8,4),(1,2,4),(1,),(3,),"reverse"),
        ((1,8,4),(1,0,4),(1,),(2,),"plain"),
        ((1,8,4),(1,2,4),(-2,),(4,),"scalar_update"),
        ((1,16,128),(1,3,128),(1,),(5,),"plain")]
    for i,(shape,size,axes,values,kind) in enumerate(profiles):
        idx_dtype=getattr(mx,("int32","int64","uint32","uint64","int8","uint8","int16","uint16","int32","uint16")[i])
        phys=shape[:-1]+(shape[-1]*2,) if kind=="stride" else shape
        if kind=="transpose":phys=(shape[0],shape[2],shape[1])
        if kind=="broadcast":phys=(1,1,shape[-1])
        a=pattern(int(np.prod(phys)),221).astype(dtype).reshape(phys)
        if kind=="stride":a=a[...,::2]
        if kind=="transpose":a=a.swapaxes(-1,-2)
        if kind=="broadcast":a=mx.broadcast_to(a,shape)
        if kind=="reverse":a=a[:,::-1,:]
        idx=mx.array(values,dtype=idx_dtype)
        if i==0:idx=idx.reshape(())
        if i==2:idx=mx.array([values[0],99],dtype=idx_dtype)[:1]
        upd=pattern(int(np.prod(size)),222).astype(dtype).reshape(size)
        if kind=="scalar_update":upd=mx.array(0.5,dtype=mx.float32)
        sliced=mx.slice(a,idx,axes=axes,slice_size=size)
        changed=mx.slice_update(a,upd,idx,axes=axes)
        # Reuse outputs as subsequent input, mirroring moving cache windows.
        back=mx.slice(changed,idx,axes=axes,slice_size=size)
        emit(f"dynamic_slice_graph_{dtype_tag}_{i}",[sliced,changed,back],dtype=dtype_tag,idx_dtype=str(idx_dtype).split('.')[-1],profile=i,shape=shape,size=size,axes=axes,values=values,kind=kind,shapes=[x.shape for x in (sliced,changed,back)])


def qsa_prefill_graph_fixtures(dtype_tag):
    """Unmodified producer-selected prefill plus standalone strided top-k."""
    from mtplx.kernels.qsa_indexer_prefill import (qsa_indexer_prefill_metal,
        qsa_indexer_prefill_scores_mpp_supported, qsa_indexer_prefill_score_chunk_rows,
        qsa_indexer_prefill_topk_metal)
    assert mx.__version__=="0.32.2"
    tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0"
    profiles=[(2,4,128,513,512,4,2051,"plain",128*1024*1024),
        (65,4,128,1024,512,4,2051,"strided",1024*4*40),
        (17,4,128,37,7,3,80,"transpose",128*1024*1024),
        (5,2,7,17,3,3,9,"broadcast",128*1024*1024),
        (9,1,8,33,7,1,15,"mixed",33*4*3),
        (3,4,128,1,1,4,0,"plain",1),
        (33,4,128,64,7,4,128,"mixed",64*4*40)]
    for i,(rows,heads,dim,blocks,topk,ratio,pos,kind,budget) in enumerate(profiles):
        dtype=getattr(mx,dtype_tag)
        pdtype=(mx.float32 if i==4 else mx.float16 if dtype!=mx.float16 else mx.bfloat16) if kind=="mixed" else dtype
        def source(shape,salt,dtype):
            count=int(np.prod(shape));a=pattern(count,salt).astype(dtype)
            if dtype==mx.float32:a=a+mx.array((np.arange(count)%7).astype(np.float32)/65536)
            return a.reshape(shape)
        qs=(1,rows,dim,heads) if kind=="transpose" else (1,1 if kind=="broadcast" else rows,heads,dim*(2 if kind=="strided" else 1))
        ps=(1,dim,blocks) if kind=="transpose" else (1,blocks,dim*(2 if kind=="strided" else 1))
        q=source(qs,211,dtype);p=source(ps,212,pdtype)
        if kind=="transpose":q=q.swapaxes(-1,-2);p=p.swapaxes(-1,-2)
        if kind=="strided":q=q[...,::2];p=p[...,::2]
        if kind=="broadcast":q=mx.broadcast_to(q,(1,rows,heads,dim))
        mpp=qsa_indexer_prefill_scores_mpp_supported(q,p)
        chunk=qsa_indexer_prefill_score_chunk_rows(rows,heads,blocks,budget,producer="mpp" if mpp else "mlx")
        total=pos+rows
        for mode in ("blocks","row_tokens","dense_mask"):
            out=qsa_indexer_prefill_metal(q,p,pos_start=pos,total_tokens=total,logical_blocks=total//ratio,block_topk=topk,compress_ratio=ratio,output_total_tokens=total+3,mode=mode,score_workspace_bytes=budget)
            outputs=[out] if mode=="dense_mask" else list(out)
            emit(f"qsa_prefill_graph_{dtype_tag}_{i}_{mode}_tf{int(tf32)}",outputs,dtype=dtype_tag,pdtype=str(pdtype).split('.')[-1],tf32=tf32,rows=rows,heads=heads,dim=dim,blocks=blocks,topk=topk,ratio=ratio,pos=pos,kind=kind,budget=budget,mpp=mpp,chunk=chunk,mode=mode,shapes=[a.shape for a in outputs])
    if dtype_tag=="float32":
        for i,(rows,blocks,topk,ratio,pos,kind) in enumerate(((1,8193,512,1,8192,"stride"),(3,2049,7,4,8190,"transpose"),(2,65,3,3,7,"broadcast"))):
            shape=(blocks,rows) if kind=="transpose" else (1 if kind=="broadcast" else rows,blocks*2 if kind=="stride" else blocks)
            scores=pattern(int(np.prod(shape)),213).astype(mx.float32).reshape(shape)
            if kind=="transpose":scores=scores.T
            if kind=="stride":scores=scores[:,::2]
            if kind=="broadcast":scores=mx.broadcast_to(scores,(rows,blocks))
            for mode in ("blocks","row_tokens","dense_mask"):
                total=pos+rows
                out=qsa_indexer_prefill_topk_metal(scores,pos_start=mx.array(pos,dtype=mx.int32),total_tokens=mx.array(total,dtype=mx.int32),logical_blocks=mx.array(total//ratio,dtype=mx.int32),block_topk=topk,compress_ratio=ratio,output_total_tokens=total+3,mode=mode)
                outputs=[out] if mode=="dense_mask" else list(out)
                emit(f"qsa_prefill_topk_graph_{i}_{mode}_tf{int(tf32)}",outputs,tf32=tf32,rows=rows,blocks=blocks,topk=topk,ratio=ratio,pos=pos,kind=kind,mode=mode,shapes=[a.shape for a in outputs])


def qsa_fused_graph_fixtures(dtype_tag):
    """Original selector factory and actual QSAIndexer chunking, no replacements."""
    from types import SimpleNamespace
    from mtplx.models.qwen4_exp import QSAIndexer
    from mtplx.kernels.qsa_indexer_select import qsa_indexer_select_metal
    assert mx.__version__=="0.32.2"
    tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0"
    profiles=[(1,4,128,513,512,4,2051,"plain"),
              (8,4,128,1024,512,4,2051,"strided"),
              (3,1,8,17,7,1,9,"transpose"),
              (5,2,7,5,3,3,7,"broadcast"),
              (2,4,64,0,1,4,0,"plain"),
              (2,4,64,5,1,4,0,"plain"),
              (129,1,8,65537,7,4,1024,"chunk")]
    for i,(rows,heads,dim,blocks,topk,ratio,pos,kind) in enumerate(profiles):
        qdtype=getattr(mx,dtype_tag)
        pdtype=mx.float32 if i==1 else (mx.float16 if i==2 else qdtype)
        def source(shape,salt,dtype):
            count=int(np.prod(shape))
            a=pattern(count,salt).astype(dtype)
            if dtype==mx.float32: a=a+mx.array((np.arange(count)%7).astype(np.float32)/65536)
            return a.reshape(shape)
        qshape=(1,rows,dim,heads) if kind=="transpose" else (1,1 if kind=="broadcast" else rows,heads,dim*(2 if kind=="strided" else 1))
        pshape=(1,dim,blocks) if kind=="transpose" else (1,blocks,dim*(2 if kind=="strided" else 1))
        q=source(qshape,201,qdtype);p=source(pshape,202,pdtype)
        if kind=="transpose": q=q.swapaxes(-1,-2);p=p.swapaxes(-1,-2)
        if kind=="strided": q=q[...,::2];p=p[...,::2]
        if kind=="broadcast": q=mx.broadcast_to(q,(1,rows,heads,dim))
        total=pos+rows;logical=total//ratio
        for mode in ("blocks","row_tokens","dense_mask"):
            if kind=="chunk":
                indexer=SimpleNamespace(block_topk=topk,ratio=ratio,_fused_score_scratch_bytes=32*1024*1024)
                # Bind the original helper, not a fixture implementation.
                from types import MethodType
                indexer._fused_query_chunk_rows=MethodType(QSAIndexer._fused_query_chunk_rows,indexer)
                out=QSAIndexer._select_fused(indexer,q,pos,p,logical,total,mode)
            else:
                out=qsa_indexer_select_metal(q,p,pos_start=pos,total_tokens=total,logical_blocks=logical,
                    block_topk=topk,compress_ratio=ratio,output_total_tokens=total+3 if mode=="dense_mask" else None,mode=mode)
            outputs=[out] if mode=="dense_mask" else list(out)
            emit(f"qsa_fused_graph_{dtype_tag}_{i}_{mode}_tf{int(tf32)}",outputs,dtype=dtype_tag,pdtype=str(pdtype).split('.')[-1],
                tf32=tf32,profile=i,rows=rows,heads=heads,dim=dim,blocks=blocks,topk=topk,ratio=ratio,pos=pos,kind=kind,mode=mode,shapes=[a.shape for a in outputs])


def qsa_indexer_prefix_fixtures(dtype_tag):
    """Unmodified QSAIndexer call, stateful prefix/rollback, actual cache."""
    import mlx.nn as nn
    from mtplx.models.qwen4_exp import QSAIndexer, QSACache
    assert mx.__version__=="0.32.2"
    tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0"
    dtype=getattr(mx,dtype_tag)
    profiles=[(128,64,False,"dense"),(128,128,True,"dense"),(64,32,True,"supplied"),
              (256,64,True,"dense"),(128,64,True,"promoted")]
    if dtype_tag=="bfloat16": profiles += [(128,64,False,"q4"),(128,64,True,"q8")]
    for profile,(dim,rot,fused,kind) in enumerate(profiles):
        # Real supported reference configuration: eager selection with either
        # original preparation branch. S>=2 is not legacy-fused eligible;
        # decode-gather excludes that selector on S1. No method replacements.
        for key,value in dict(MTPLX_COMPILED_QSA_INDEXER=0,MTPLX_FUSED_QSA_INDEXER=int(fused),
             MTPLX_QSA_PREFILL=0,MTPLX_QSA_PREFILL_MIN_ROWS=2,MTPLX_QSA_GATHER=1,
             MTPLX_QSA_GATHER_MIN_CONTEXT=0,MTPLX_QSA_GATHER_MAX_ROWS=8,
             MTPLX_QSA_GATHER_DECODE=1,MTPLX_QSA_FLASH=0,MTPLX_QSA_SCORE_TILE_ROWS=0).items(): os.environ[key]=str(value)
        indexer=QSAIndexer.__new__(QSAIndexer)
        nn.Module.__init__(indexer)
        indexer.n_heads=4;indexer.kv_heads=1;indexer.head_dim=dim;indexer.ratio=4
        indexer.block_topk=512;indexer.budget=2048;indexer.rms_norm_eps=1e-6
        indexer._inv_freq=mx.array([1/(j+1) for j in range(rot//2)],dtype=mx.float32)
        indexer._rope_attention_scaling=1.125
        indexer.q_layernorm=nn.RMSNorm(dim,eps=1e-6)
        indexer.k_layernorm=nn.RMSNorm(dim,eps=1e-6)
        norm_dtype=mx.float32 if kind=="promoted" else dtype
        indexer.q_layernorm.weight=pattern(dim,185).astype(norm_dtype)
        indexer.k_layernorm.weight=pattern(dim,186).astype(norm_dtype)
        if kind in ("q4","q8"):
            indexer.index_qk_proj=quantized_linear_fixture(5*dim,64,int(kind[1:]),64,187)
        else:
            indexer.index_qk_proj=nn.Linear(64,5*dim,bias=False)
            indexer.index_qk_proj.weight=pattern(5*dim*64,187).astype(dtype).reshape(5*dim,64)
        cache=QSACache()
        cache.reserve_indexer_capacity(raw_capacity=256,pooled_capacity=256)
        for step,(rows,trim) in enumerate(((3,0),(1,0),(5,0),(2041,0),(1,0),(4,0),(8,5),(33,0),(1,0))):
            if trim: cache.trim(trim)
            pos=cache.offset
            hidden=pattern(rows*64,190+step*7).astype(dtype).reshape(1,rows,64)
            qk=None
            if kind=="supplied":
                qk=pattern(rows*5*dim*2,191+step*7).astype(dtype).reshape(1,rows,5*dim*2)[...,::2]
            result=indexer(hidden,pos,cache,qk_rows=qk)
            outputs=[cache.raw_keys]
            names=["raw"]
            if cache.pooled is not None: outputs += [cache.pooled,cache.pooled_f32_t];names += ["pooled","mirror"]
            if result is None: lane="none"
            elif isinstance(result,tuple): lane=result[0];outputs+=list(result[1:]);names+=["indices","valid"]
            else: lane="decode_tokens" if result.ndim==1 else "dense";outputs.append(result);names.append("selection")
            emit(f"qsa_indexer_prefix_{dtype_tag}_{profile}_{step}_tf{int(tf32)}",outputs,
                 dtype=dtype_tag,tf32=tf32,profile=profile,dim=dim,rot=rot,fused=fused,kind=kind,
                 step=step,rows=rows,trim=trim,pos=pos,offset=cache.offset,pooled_len=cache.pooled_len,
                 lane=lane,names=names,shapes=[a.shape for a in outputs])
            assert cache.offset==pos  # Only Attention advances this frontier.
            cache.kv.offset=pos+rows


def qsa_eager_graph_fixtures(dtype_tag):
    """Actual QSAIndexer._select_eager with its actual QSACache and all exits."""
    from types import SimpleNamespace, MethodType
    from mtplx.models.qwen4_exp import QSAIndexer, QSACache, _qsa_large_prefill_enabled
    from mtplx.attention_context import attention_phase
    assert mx.__version__=="0.32.2"
    tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0"
    profiles=[(1,2051,4,512,0,m) for m in ("dense","flash","decode","all")]
    profiles += [(8,2052,4,512,t,m) for t,m in ((0,"dense"),(0,"rows"),(3,"rows"),(8,"rows"),(9,"rows"),(0,"min"),(0,"max"))]
    profiles += [(137,2052,4,512,t,m) for t,m in ((0,"prefill"),(31,"prefill"),(0,"all"),(31,"all"),(0,"verify"))]
    profiles += [(2048,4,4,512,0,"dense"),(2,32768,4,512,0,"prefill"),
                 (33,32768,4,512,8,"prefill"),(1,4096,4,512,0,"decode"),
                 (8,2052,8,512,0,"prefill"),(3,13,1,7,0,"rows"),
                 (3,13,1,7,1,"dense"),(1,13,1,7,0,"flash")]
    for i,(rows,pos,ratio,topk,tile,mode) in enumerate(profiles):
        flash=mode in ("flash","all")
        decode=mode in ("decode","all")
        gather=mode in ("rows","all","min","max")
        prefill=mode in ("prefill","all","verify")
        phase="decode_verify" if mode=="verify" else "prefill"
        minimum=16384 if mode=="min" else 0
        maximum=7 if mode=="max" else 8
        for key,value in dict(MTPLX_QSA_FLASH=int(flash),MTPLX_QSA_GATHER_DECODE=int(decode),
              MTPLX_QSA_GATHER=int(gather),MTPLX_QSA_GATHER_MIN_CONTEXT=minimum,MTPLX_QSA_GATHER_MAX_ROWS=maximum,
              MTPLX_QSA_PREFILL=int(prefill),MTPLX_QSA_PREFILL_MIN_ROWS=2,MTPLX_QSA_PREFILL_MIN_CONTEXT=2049,
              MTPLX_QSA_SCORE_TILE_ROWS=tile).items(): os.environ[key]=str(value)
        total=pos+rows
        nb=total//ratio
        dtype=getattr(mx,dtype_tag)
        q=pattern(rows*4*128,181).astype(dtype)
        if dtype==mx.float32: q=q+mx.array((np.arange(q.size)%7).astype(np.float32)/65536)
        q=q.reshape(1,rows,4,128)
        pooled=pattern(nb*128,182).astype(dtype).reshape(1,nb,128)
        cache=QSACache(compress_ratio=ratio)
        cache.write_pooled(pooled,0,nb)
        indexer=SimpleNamespace(head_dim=128,ratio=ratio,block_topk=topk)
        indexer._tiled_topk=MethodType(QSAIndexer._tiled_topk,indexer)
        with attention_phase(phase):
            large=_qsa_large_prefill_enabled(rows,total)
            result=QSAIndexer._select_eager(indexer,q,pos,cache,pooled,total)
        tail=None
        if isinstance(result,tuple):
            lane=result[0]
            if lane=="flash": outputs=[result[1]];tail=result[2]
            else: outputs=list(result[1:])
        else:
            lane="decode_tokens" if result.ndim==1 else "dense"
            outputs=[result]
        emit(f"qsa_eager_graph_{dtype_tag}_{i}_tf{int(tf32)}",outputs,dtype=dtype_tag,tf32=tf32,
             rows=rows,pos=pos,ratio=ratio,topk=topk,tile=tile,mode=mode,large=large,flash=flash,decode=decode,
             gather=gather,minimum=minimum,maximum=maximum,lane=lane,tail=tail,shapes=[a.shape for a in outputs])


def sort_graph_fixtures(dtype_tag):
    """All four original sort primitives, axis/layout and merge boundaries."""
    assert mx.__version__ == "0.32.2"
    dtype=getattr(mx,dtype_tag)
    profiles=[((2,w),-1,"plain") for w in (1,128,129,256,257,512,513,1024,1025,2048,2049,4097,8193)]
    profiles += [((3,7,5),axis,kind) for axis,kind in ((0,"plain"),(1,"plain"),(2,"transpose"),(1,"strided"),(-1,"reverse"),(1,"broadcast"),(None,"transpose"))]
    profiles += [((2049,3),0,"plain"),((2,2049,3),1,"strided"),((4097,),0,"offset")]
    for i,(shape,axis,kind) in enumerate(profiles):
        physical=list(shape)
        if kind=="strided": physical[-1]*=2
        if kind=="broadcast": physical[1]=1
        count=int(np.prod(physical))+(kind=="offset")
        values=((np.arange(count,dtype=np.int64)*17+31)%257)
        if dtype_tag not in ("uint32","uint64"): values-=128
        a=mx.array(values).astype(dtype)
        if kind=="offset": a=a[1:]
        a=a.reshape(physical)
        if kind=="transpose": a=a.swapaxes(0,-1)
        if kind=="strided": a=a[...,::2]
        if kind=="reverse": a=a[...,::-1]
        if kind=="broadcast": a=mx.broadcast_to(a,shape)
        width=a.size if axis is None else a.shape[axis]
        kth=-min(3,width)
        outputs=[mx.sort(a,axis=axis),mx.argsort(a,axis=axis),mx.partition(a,kth,axis=axis),mx.argpartition(a,kth,axis=axis)]
        emit(f"sort_graph_{dtype_tag}_{i}",outputs,dtype=dtype_tag,profile=i,shape=shape,axis=axis,kind=kind,kth=kth,output_shape=outputs[0].shape)


def dense_float_fixtures(dtype_tag=None, start=0, end=None):
    """Pinned Matmul factory/backend including scalar, vector and empty cases."""
    assert mx.__version__ == "0.32.2"
    tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0"
    profiles=[(m,n,k,"plain",ta,tb) for m,n,k in ((3,5,128),(16,17,129),(40,40,127),(40,40,128))
              for ta in (False,True) for tb in (False,True)]
    profiles += [(m,n,k,kind,ta,tb) for m,n,k,kind,ta,tb in (
        (1,1,7,"vectors",False,False),(1,1,16385,"vectors",False,False),
        (1,17,64,"left-vector",False,False),(1,17,65,"left-vector",False,True),
        (7,1,65,"right-vector",False,False),(7,1,65,"right-vector",True,False),
        (1,1,129,"cross",False,False),(7,1,129,"cross",True,False),
        (5,7,64,"cross",False,True),(5,7,64,"copy",False,False),
        (5,7,64,"offset",False,True),(5,7,64,"collapse",False,True),
        (5,7,64,"flatten",False,True),(65,513,128,"plain",False,True),
        (2,2,4097,"plain",False,False),(193,17,8193,"plain",False,True),
        (0,5,7,"plain",False,False),(5,0,7,"plain",False,False),
        (5,7,0,"plain",False,False),(0,5,0,"flatten",False,False),
        (3,7,128,"mixed",False,True),(3,7,128,"mixed-f32",False,True),
    )]
    for dtype in (mx.bfloat16,mx.float16,mx.float32):
        tag=str(dtype).split('.')[-1]
        if dtype_tag is not None and tag!=dtype_tag: continue
        for i,(m,n,k,kind,ta,tb) in enumerate(profiles):
            if i<start or (end is not None and i>=end): continue
            def source(shape,salt,typ):
                physical=list(shape)
                if kind=="copy": physical[-1]*=2
                count=int(np.prod(physical))+(kind=="offset")
                a=pattern(count,salt).astype(typ)
                if typ==mx.float32: a=a+mx.array((np.arange(count)%7).astype(np.float32)/65536)
                if kind=="offset": a=a[1:]
                a=a.reshape(physical)
                return a[...,::2] if kind=="copy" else a
            abatch=(2,1) if kind=="cross" else ((3,) if kind in ("copy","offset","collapse","flatten") else ())
            bbatch=(1,3) if kind=="cross" else ((3,) if kind in ("copy","offset") else ((1,) if kind=="collapse" else ()))
            adtype=mx.bfloat16 if kind.startswith("mixed") else dtype
            bdtype=(mx.float32 if kind=="mixed-f32" else mx.float16) if kind.startswith("mixed") else dtype
            a=source(abatch+((k,m) if ta else (m,k)),171,adtype)
            b=source(bbatch+((n,k) if tb else (k,n)),172,bdtype)
            if ta: a=a.swapaxes(-2,-1)
            if tb: b=b.swapaxes(-2,-1)
            if kind in ("left-vector","vectors"): a=a.reshape(k)
            if kind in ("right-vector","vectors"): b=b.reshape(k)
            out=a@b
            emit(f"dense_float_{tag}_{i}_tf{int(tf32)}",[out],dtype=tag,profile=i,m=m,n=n,k=k,kind=kind,ta=ta,tb=tb,tf32=tf32,
                 a_dtype=str(adtype).split('.')[-1],b_dtype=str(bdtype).split('.')[-1],output_shape=out.shape)


def sum_graph_fixtures(dtype_tag=None, start=0, end=None):
    """Original runtime reductions used by QSA/HC/PLE/GDN; no custom sums."""
    assert mx.__version__ == "0.32.2"
    profiles = [
        ((4096,), (0,), "plain"), ((4097,), (0,), "offset"),
        ((31,128), (-1,), "plain"), ((32,128), (-1,), "plain"),
        ((32,512), (1,), "plain"), ((32,513), (1,), "plain"),
        ((32,1024), (1,), "plain"), ((32,1025), (1,), "plain"),
        ((2,4097), (1,), "plain"), ((5,7), (1,), "plain"),
        ((1,1024,4,128), (2,), "plain"), ((1,137,4,513), (2,), "plain"),
        ((31,128), (0,), "plain"), ((32,128), (0,), "plain"),
        ((256,128), (0,), "plain"), ((257,128), (0,), "plain"),
        ((1023,17), (0,), "plain"), ((1024,17), (0,), "plain"),
        ((32768,2), (0,), "plain"),
        ((3,5,7,128), (0,2), "plain"),
        ((3,5,7,9,2,128), (0,2,4), "plain"),
        ((3,5,7,128), (0,3), "plain"),
        ((3,5,7,9,2,128), (0,2,5), "plain"),
        ((64,7,4), (0,2), "plain"),
        ((64,7,16), (0,2), "plain"),
        ((5,7,128), (-1,), "transpose"),
        ((5,7,128), (1,), "transpose"),
        ((5,7,128), (0,1,2), "transpose"),
        ((5,7,128), (-1,), "strided"),
        ((5,7,128), (0,), "strided"),
        ((5,7,128), (1,), "reverse"),
        ((5,7,128), (-1,), "broadcast"),
        ((5,7,128), (0,), "broadcast"),
        ((5,7,128), (0,1,2), "scalar"),
        ((1,7,1), (0,2), "offset"),
        ((0,7,128), (0,), "plain"),
        ((5,0,128), (-1,), "plain"),
        ((5,7,128), (), "plain"),
        ((257,32768), (0,), "plain"),
        ((16777217,), (0,), "plain"),
    ]
    for dtype in (mx.bfloat16,mx.float16,mx.float32):
        tag=str(dtype).split('.')[-1]
        if dtype_tag is not None and tag != dtype_tag: continue
        for i,(shape,axes,layout) in enumerate(profiles):
            if i<start or (end is not None and i>=end): continue
            if i>=38 and dtype!=mx.float32: continue
            base=list(shape)
            if layout=="broadcast": base[0]=1
            if layout=="strided": base[-1]*=2
            n=int(np.prod(base))
            a=pattern(1 if layout=="scalar" else n+(layout=="offset"),170).astype(dtype)
            if layout=="offset": a=a[1:]
            a=mx.broadcast_to(a.reshape(()),shape) if layout=="scalar" else a.reshape(base)
            if layout=="strided": a=a[...,::2]
            if layout=="reverse": a=a[...,::-1]
            if layout=="broadcast": a=mx.broadcast_to(a,shape)
            if layout=="transpose": a=a.swapaxes(0,1)
            for keep in (False,True):
                out=mx.sum(a,axis=axes,keepdims=keep)
                mean=mx.mean(a.astype(mx.float32),axis=axes,keepdims=keep)
                emit(f"sum_graph_{tag}_{i}_{keep}",[out,mean],dtype=tag,profile=i,shape=shape,
                     axes=axes,layout=layout,keep=keep,output_shape=out.shape)


def qsa_dense_mask_fixtures():
    from mtplx.models.qwen4_exp import _qsa_blocks_to_dense_mask
    assert mx.__version__ == "0.32.2"
    for rows,pos,ratio,topk in ((1,0,4,0),(1,0,4,8),(2,1,4,8),(5,0,4,9),
                               (8,3,4,17),(9,2045,4,512),(137,32767,4,512),
                               (2048,2048,4,512),(5,2,1,9),(8,17,7,8),(0,0,4,0)):
        total=pos+rows
        for strided in (False,True):
            for empty in (False,True):
                cols=topk*(2 if strided else 1)
                ids=[];valid=[]
                for r in range(rows):
                    for c in range(cols):
                        slot=c//2 if strided else c
                        candidate=(pos+r+1)//ratio-1-slot//2
                        ids.append(total if slot%11==5 else (-1 if slot%11==3 else candidate))
                        valid.append(not empty and slot%7!=2)
                ids=mx.array(ids,dtype=mx.int32).reshape(rows,cols)
                valid=mx.array(valid,dtype=mx.bool_).reshape(rows,cols)
                if strided: ids,valid=ids[:,::2],valid[:,::2]
                out=_qsa_blocks_to_dense_mask(ids,valid,pos_start=pos,total_tokens=total,compress_ratio=ratio)
                emit(f"qsa_dense_mask_{rows}_{pos}_{ratio}_{topk}_{strided}_{empty}",[out],
                     rows=rows,pos=pos,ratio=ratio,topk=topk,strided=strided,empty=empty)


def scatter_axis_graph_fixtures():
    assert mx.__version__ == "0.32.2"
    for shape,axis in (((7,),0),((3,5),0),((3,5),1),((2,3,5),1)):
        for strided in (False,True):
            for scalar in (False,True):
                for src_layout in ("plain","scalar","offset","strided"):
                    n=int(np.prod(shape))
                    if src_layout=="scalar": src=mx.broadcast_to(mx.array(False),shape)
                    else:
                        mult=2 if src_layout=="strided" else 1
                        src=(mx.arange(n*mult+1,dtype=mx.int32)%3==0)
                        src=src[1:] if src_layout=="offset" else src[:n*mult]
                        if src_layout=="strided": src=src[::2]
                        src=src.reshape(shape)
                    ishape=list(shape);ishape[axis]=2
                    size=int(np.prod(ishape))
                    mult=2 if strided else 1
                    idx=mx.arange(size*mult,dtype=mx.int64)//mult%2
                    idx=mx.where(idx==0,-1,0)
                    if strided: idx=idx[::2]
                    idx=idx.reshape(ishape)
                    upd=mx.array(True) if scalar else (mx.arange(size,dtype=mx.int32)%3==0).reshape(ishape)
                    out=mx.put_along_axis(src,idx,upd,axis=axis)
                    emit(f"scatter_axis_{shape}_{axis}_{strided}_{scalar}_{src_layout}",[out],
                         shape=shape,axis=axis,strided=strided,scalar=scalar,src_layout=src_layout)


def sparse_attention_graph_fixtures():
    """Actual MTPLX sparse attention consumers, never a dense substitute."""
    from mtplx.kernels.qsa_flash_skip import qsa_flash_skip
    from mtplx.kernels.qsa_prefill_flash import qsa_prefill_flash, qsa_prefill_flash_supported
    assert mx.__version__ == "0.32.2"
    for dtype in (mx.bfloat16,mx.float16):
        tag=str(dtype).split('.')[-1]
        for heads,total,nsel in ((24,3,0),(24,4099,512),(2,2052,17)):
            for layout in ("plain","offset","strided"):
                cap=total+7
                def source(shape,salt):
                    n=int(np.prod(shape))
                    if layout=="offset": return pattern(n+1,salt).astype(dtype)[1:].reshape(shape)
                    if layout=="strided": return pattern(n*2,salt).astype(dtype).reshape(*shape[:-1],shape[-1]*2)[...,::2]
                    return pattern(n,salt).astype(dtype).reshape(shape)
                q=source((heads,256),111)
                k,v=source((1,2,cap,256),112),source((1,2,cap,256),113)
                ids=mx.arange(nsel*2 if layout=="strided" else nsel,dtype=mx.int32)
                if layout=="strided": ids=ids[::2]
                else: ids=ids*2
                tail=(total//4)*4
                out=qsa_flash_skip(q,k,v,ids,total,tail,0.0625)
                emit(f"sparse_skip_{tag}_{heads}_{total}_{layout}",[out],kind="skip",dtype=tag,heads=heads,total=total,nsel=nsel,layout=layout)
        for rows,pos in ((2,2050),(8,8191),(137,32767),(2048,2048)):
            for layout in ("plain","offset","strided","holes"):
                total=pos+rows;cap=total+7
                def source(shape,salt):
                    n=int(np.prod(shape))
                    if layout=="offset": return pattern(n+1,salt).astype(dtype)[1:].reshape(shape)
                    if layout=="strided": return pattern(n*2,salt).astype(dtype).reshape(*shape[:-1],shape[-1]*2)[...,::2]
                    return pattern(n,salt).astype(dtype).reshape(shape)
                q=source((1,rows,24,256),114).transpose(0,2,1,3)
                k,v=source((1,2,cap,256),115),source((1,2,cap,256),116)
                slots=mx.arange(512,dtype=mx.int32)[None,:]
                ids=mx.broadcast_to(slots,(rows,512))
                valid=mx.broadcast_to((slots<17) | ((slots==511) if layout=="holes" else mx.array(False)),(rows,512))
                if layout=="holes":
                    ids=mx.where(slots==3,-1,mx.where(slots==5,total,ids))
                if layout=="strided":
                    ids=mx.stack([ids,ids],axis=-1).reshape(rows,1024)[:,::2]
                    valid=mx.stack([valid,valid],axis=-1).reshape(rows,1024)[:,::2]
                assert qsa_prefill_flash_supported(q,k,v,ids,valid,pos_start=pos,total_tokens=total,scale=0.0625)
                out=qsa_prefill_flash(q,k,v,ids,valid,pos_start=pos,total_tokens=total,scale=0.0625)
                emit(f"sparse_prefill_{tag}_{rows}_{pos}_{layout}",[out],kind="prefill",dtype=tag,rows=rows,pos=pos,total=total,layout=layout)


def qsa_prepare_graph_fixtures(*, dtype_tag=None, profile_only=None):
    """Actual parameterized MTPLX query/pool kernels, including traced starts."""
    from mtplx.kernels.qsa_indexer_prepare import (
        qsa_indexer_prepare_queries_metal,qsa_indexer_pool_keys_metal,qsa_indexer_prepare_supported)
    assert mx.__version__ == "0.32.2"
    profiles=((1,4,128,64,4,1e-6,1.0,0),(3,1,2,2,1,1e-4,1.25,-7),
              (7,3,17,16,7,0.0,0.5,2047),(137,7,64,32,3,1e-5,1.1,65535),
              (2048,2,128,128,8,1e-6,1.3,131071),(1,4,128,64,4,0.25,1.0,0))
    for dtype in (mx.bfloat16,mx.float16,mx.float32):
        tag=str(dtype).split('.')[-1]
        if dtype_tag is not None and tag!=dtype_tag: continue
        for profile,(rows,heads,dim,rotary,ratio,eps,scale,pos) in enumerate(profiles):
            if profile_only is not None and profile!=profile_only: continue
            for layout in ("plain","offset","strided","broadcast"):
                def source(shape,salt,ty=dtype):
                    n=int(np.prod(shape))
                    if layout=="broadcast": return mx.broadcast_to(pattern(1,salt).astype(ty),shape)
                    if layout=="offset": return pattern(n+1,salt).astype(ty)[1:].reshape(shape)
                    if layout=="strided": return pattern(n*2,salt).astype(ty).reshape(*shape[:-1],shape[-1]*2)[...,::2]
                    return pattern(n,salt).astype(ty).reshape(shape)
                q=source((1,rows,heads,dim),160)
                keys=source((1,rows*ratio,dim),161)
                norm=source((dim,),162)
                freq=source((rotary//2,),163,mx.float32)
                start=pos
                if layout=="offset": start=mx.array([123,pos],dtype=mx.int32)[1:].reshape(1,1)
                elif layout=="strided": start=mx.array([[pos+5]],dtype=mx.int32)-5
                elif layout=="broadcast": start=mx.array(pos,dtype=mx.int32)
                assert qsa_indexer_prepare_supported(q,norm,freq,expected_ndim=4)
                assert qsa_indexer_prepare_supported(keys,norm,freq,expected_ndim=3)
                out=qsa_indexer_prepare_queries_metal(q,norm,freq,pos_start=start,eps=eps,attention_scaling=scale)
                pooled=qsa_indexer_pool_keys_metal(keys,norm,freq,block_start=start,compress_ratio=ratio,eps=eps,attention_scaling=scale)
                emit(f"qsa_prepare_graph_{tag}_{profile}_{layout}",[out,pooled],dtype=tag,profile=profile,
                     layout=layout,rows=rows,heads=heads,dim=dim,rotary=rotary,ratio=ratio,eps=eps,scale=scale,pos=pos)


def qsa_cache_graph_fixtures(*, profile=None, deferred_only=None):
    """Actual positional QSACache methods, retaining Python array aliases."""
    from types import SimpleNamespace
    from mtplx.models.qwen4_exp import QSACache,QSAIndexer
    assert mx.__version__ == "0.32.2"
    profiles=[(n,4,128,mx.bfloat16) for n in (0,1,255,256,257,1025)]
    profiles += [(257,1,17,mx.float16),(257,7,32,mx.float32)]
    for initial,ratio,dim,dtype in profiles:
        if profile is not None and (initial,ratio)!=tuple(profile): continue
        for deferred in (False,True):
            if deferred_only is not None and deferred!=deferred_only: continue
            cache=QSACache(ratio)
            outputs=[];snapshots=[];metadata=[]
            tag=str(dtype).split('.')[-1]
            actions=[("reserve",768,320),("append",initial,0),("snapshot",0,0),
                     ("reserve",2049,769),("view",0,0),("append",3,0),
                     ("trim",2,0),("append",4,0),("restore",0,0),("view",0,0),
                     ("append",257,0),("trim",999999,0),("append",1,0),
                     ("ensure",1025,0),("append",4,0)]
            def source(shape,salt,strided):
                n=int(np.prod(shape));mult=2 if strided else 1
                a=pattern(n*mult,salt).astype(dtype).reshape(*shape[:-1],shape[-1]*mult)
                return a[...,::2] if strided else a
            def append(n,step):
                start=cache.offset;total=start+n
                cache.write_raw(source((1,n,dim),140+step*3,step%2==1))
                old=min(cache.pooled_len,total//ratio)
                if total//ratio>old:
                    cache.write_pooled(source((1,total//ratio-old,dim),141+step*3,step%2==1),old,total//ratio)
                cache.kv.update_and_fetch(
                    source((1,n,2,256),142+step*3,False).transpose(0,2,1,3),
                    source((1,n,2,256),143+step*3,False).transpose(0,2,1,3))
            for step,(action,a,b) in enumerate(actions):
                extra={}
                if action=="reserve": cache.reserve_indexer_capacity(raw_capacity=a,pooled_capacity=b)
                elif action=="append": append(a,step)
                elif action=="snapshot": snapshots.append(cache.state)
                elif action=="trim": extra["trimmed"]=cache.trim(a)
                elif action=="restore": cache.state=snapshots[a]
                elif action=="view" and cache.pooled is not None:
                    outputs.append(cache.pooled_f32_view(cache.pooled_len))
                elif action=="ensure":
                    indexer=SimpleNamespace(ratio=ratio,head_dim=dim)
                    raw,pooled=QSAIndexer._ensure_compiled_backings(indexer,cache,dtype=dtype,pos_start=cache.offset,rows=a)
                    assert raw is cache.raw_keys and pooled is cache.pooled
                current=[cache.kv.keys,cache.kv.values,cache.raw_keys,cache.pooled,cache.pooled_f32_t]
                if cache.kv.keys is not None: current.extend(cache.state)
                outputs.extend(v for v in current if v is not None)
                if not deferred: mx.eval(*(v for v in outputs if v is not None))
                metadata.append(dict(action=action,a=a,b=b,offset=cache.offset,pooled_len=cache.pooled_len,
                    capacities=[0 if v is None else v.shape[axis] for v,axis in zip(current[:5],(2,2,1,1,3))],
                    reserved=[cache._reserved_raw_capacity,cache._reserved_pooled_capacity],nbytes=cache.nbytes,**extra))
            emit(f"qsa_cache_graph_{initial}_{ratio}_{dim}_{tag}_{int(deferred)}",outputs,
                 initial=initial,ratio=ratio,dim=dim,dtype=tag,deferred=deferred,actions=metadata,
                 shapes=[list(v.shape) for v in outputs])


def kv_cache_graph_fixtures():
    """Actual KVCache mutations, retaining old states across lazy updates."""
    from mlx_lm.models.cache import KVCache
    for initial in (0,1,255,256,257,2048):
        for deferred in (False,True):
            cache=KVCache()
            outputs=[]
            snapshots=[]
            metadata=[]
            actions=[("append",initial),("append",3),("trim",2),("append",4),
                     ("restore",0),("append",257),("trim",999999),("append",1)]
            for step,(action,n) in enumerate(actions):
                if action=="trim":
                    trimmed=cache.trim(n)
                    metadata.append(dict(action=action,n=n,trimmed=trimmed,offset=cache.offset))
                    continue
                if action=="restore":
                    cache.state=snapshots[0]
                    metadata.append(dict(action=action,n=n,offset=cache.offset))
                    continue
                k=pattern(n*2*256,104+step*2).reshape(1,n,2,256).transpose(0,2,1,3)
                v=pattern(n*2*256,105+step*2).reshape(1,n,2,256).transpose(0,2,1,3)
                state=cache.update_and_fetch(k,v)
                snapshots.append(state)
                current=[*state,cache.keys,cache.values]
                outputs.extend(current)
                if not deferred: mx.eval(*current)
                metadata.append(dict(action=action,n=n,offset=cache.offset,capacity=cache.keys.shape[2]))
            emit(f"kv_cache_graph_{initial}_{int(deferred)}",outputs,initial=initial,deferred=deferred,actions=metadata)


def rope_graph_fixtures():
    """Original Qwen RoPE functions, including static YaRN and vision axes."""
    from types import SimpleNamespace
    from mtplx.models.qwen4_exp import (_rope_inv_freq_and_scaling, _rope_cos_sin,
        _mrope_cos_sin, _build_mrope_axes, _apply_partial_rope)
    assert mx.__version__ == "0.32.2"
    profiles = [
        (64, 10000000., 1, 0, {}),
        (256, 10000., 7, 511, {}),
        (64, 10000000., 2048, 262143, dict(rope_type="yarn", factor=4., original_max_position_embeddings=262144)),
        (64, 10000., 7, 1048576, dict(rope_type="yarn", factor=1., original_max_position_embeddings=128)),
        (64, 10000., 64, 8191, dict(rope_type="yarn", factor=8., original_max_position_embeddings=4096,
            mscale=2., mscale_all_dim=1., truncate=False, beta_fast=4., beta_slow=4.)),
        (64, 10000., 7, 32767, dict(rope_type="yarn", factor=4., original_max_position_embeddings=128,
            attention_factor=0.75, beta_fast=0., beta_slow=0.)),
    ]
    for profile,(rot,base,rows,pos,parameters) in enumerate(profiles):
        inv,scale = _rope_inv_freq_and_scaling(SimpleNamespace(rotary_dim=rot, rope_theta=base, rope_parameters=parameters))
        heads=24 if rows<=7 else 2
        half=rot//2
        section=[(half+2)//3,(half+1)//3,half//3]
        for mode in ("text","equal","interleaved","sections"):
            axes=_build_mrope_axes(section,mode!="sections")
            positions=mx.arange(pos,pos+rows,dtype=mx.int32)
            positions3=mx.stack([positions,positions+(0 if mode=="equal" else 7),positions+(0 if mode=="equal" else 19)])
            cos,sin=(_rope_cos_sin(positions,inv,scale) if mode=="text" else
                     _mrope_cos_sin(positions3,inv,mx.array(axes,mx.int32)))
            for strided in (False,True):
                width=512 if strided else 256
                x=pattern(rows*heads*width,103).reshape(1,rows,heads,width)
                if strided: x=x[...,1::2]
                out=_apply_partial_rope(x,cos,sin)
                emit(f"rope_graph_{profile}_{mode}_{int(strided)}",[inv,cos,sin,out],
                     rot=rot,base=base,rows=rows,pos=pos,parameters=parameters,heads=heads,
                     mode=mode,strided=strided,section=section,axes=axes,scale=scale)


def qsa_rows_attention_fixtures():
    """Actual MTPLX per-row gathered attention, not a reconstructed formula."""
    from mtplx.models.qwen4_exp import _qsa_rows_gather_attention
    assert mx.__version__ == "0.32.2"
    for heads,rows,total,count in ((24,1,33,8),(24,3,1057,1025),(24,8,4193,4097),
                                  (24,64,113,17),(2,3,33,8),(2,16,1057,1025)):
        for layout in ("plain","offset","strided"):
            q=pattern(rows*heads*256,97).reshape(1,rows,heads,256).transpose(0,2,1,3)
            def cached(salt):
                cap=total+7 if layout=="offset" else total*2 if layout=="strided" else total
                a=pattern(2*cap*256,salt).reshape(1,2,cap,256)
                return a[:,:,2:total+2] if layout=="offset" else a[:,:,::2] if layout=="strided" else a
            k,v=cached(98),cached(99)
            idx=(mx.arange(rows*count,dtype=mx.int32)*7%total).reshape(rows,count)
            ok=(mx.arange(rows*count,dtype=mx.int32).reshape(rows,count)%5)!=1
            emit(f"qsa_rows_attention_h{heads}_s{rows}_t{total}_k{count}_{layout}",
                 [_qsa_rows_gather_attention(q,k,v,idx,ok,0.0625)],
                 heads=heads,rows=rows,total=total,count=count,layout=layout)


def qsa_prefill_gather_fixtures():
    """Actual portable prefill function, with its explicit per-tile evals."""
    from mtplx.models.qwen4_exp import _qsa_prefill_gather_attention
    assert mx.__version__ == "0.32.2"
    for heads,rows,pos,ratio,topk in ((24,1,0,4,2),(24,7,3,4,2),(24,17,127,128,2),
                                   (24,17,511,128,2),(2,7,0,1,2),(2,2048,0,4,2)):
        for tile in (1,4,64) if rows<2048 else (64,2048):
            for strided in (False,True):
                cap=pos+rows+7
                q=pattern(rows*heads*256,100).reshape(1,rows,heads,256).transpose(0,2,1,3)
                k=pattern(2*cap*256,101).reshape(1,2,cap,256)
                v=pattern(2*cap*256,102).reshape(1,2,cap,256)
                nb=(pos+np.arange(rows)+1)//ratio
                raw_ids=nb[:,None]-1-np.arange(topk)[None,:]
                raw_ok=raw_ids>=0
                raw_ids=np.where(raw_ok,raw_ids,-999)
                if strided:
                    ids=np.full((rows,topk*2),-888,dtype=np.int64);ids[:,1::2]=raw_ids
                    ok=np.zeros((rows,topk*2),dtype=np.bool_);ok[:,1::2]=raw_ok
                    ids=mx.array(ids)[:,1::2];ok=mx.array(ok)[:,1::2]
                else:
                    ids=mx.array(raw_ids.astype(np.int32));ok=mx.array(raw_ok)
                calls=[]; original_eval=mx.eval
                def tracked(*a):
                    calls.append([list(x.shape) for x in a]);return original_eval(*a)
                mx.eval=tracked
                try:
                    out=_qsa_prefill_gather_attention(q,k,v,ids,ok,pos_start=pos,total_tokens=pos+rows,
                        compress_ratio=ratio,scale=0.0625,tile_rows=tile)
                finally:
                    mx.eval=original_eval
                emit(f"qsa_prefill_gather_h{heads}_s{rows}_p{pos}_r{ratio}_t{tile}_strided{int(strided)}",[out],
                    heads=heads,rows=rows,pos=pos,ratio=ratio,topk=topk,tile=tile,strided=strided,evals=calls)


def sdpa_vector_fixtures():
    """Actual MTPLX runtime fused D256 SDPA, including the 2-pass boundaries."""
    assert mx.__version__ == "0.32.2"
    assert not os.environ.get("MLX_SDPA_BLOCKS")
    for batch,heads,rows,total in ((1,24,1,17),(1,24,2,1023),(1,24,2,1024),
                                   (1,24,2,16384),(1,24,1,65536),(2,24,2,33)):
        for copied in (False,True):
            for mode in ("none","bool","add","causal"):
                step=2 if copied else 1
                q=pattern(batch*rows*heads*256*step,103).reshape(batch,rows,heads,256*step)[...,::step].transpose(0,2,1,3)
                k=pattern(batch*2*(total+7)*256*step,104).reshape(batch,2,total+7,256*step)[:,:,2:total+2,::step]
                v=pattern(batch*2*(total+7)*256*step,105).reshape(batch,2,total+7,256*step)[:,:,2:total+2,::step]
                raw=np.arange(rows*total).reshape(rows,total)%3!=1
                mask=mx.array(raw)[None,None] if mode=="bool" else mx.array(np.where(raw,0,-10),dtype=mx.bfloat16)[None,None] if mode=="add" else "causal" if mode=="causal" else None
                emit(f"sdpa_vector_b{batch}_h{heads}_s{rows}_t{total}_copy{int(copied)}_{mode}",
                    [mx.fast.scaled_dot_product_attention(q,k,v,scale=0.0625,mask=mask)],
                    batch=batch,heads=heads,rows=rows,total=total,copied=copied,mode=mode)


def gather_array_fixtures():
    """Pinned GatherQMM on the actual expert/broadcast/sliced Array layouts."""
    assert mx.__version__ == "0.32.2"
    # x batches, rows per matrix, output width, input width, experts, routes,
    # layout, explicit lhs, explicit rhs, sorted, route rank.
    cases = [
        (1,1,80,512,8,10,"plain",False,True,False,1),
        (4,1,80,512,8,40,"expert_slice",False,True,False,2),
        (7,1,81,64,8,70,"expert_slice",False,True,False,2),
        (2,33,80,512,8,8,"plain",True,True,False,1),
        (2,17,81,64,8,8,"transpose",True,True,False,1),
        (2,33,80,512,8,8,"expert_slice",True,True,False,1),
        (16,1,80,512,4,16,"plain",False,True,True,1),
        (256,1,80,512,4,256,"expert_slice",False,True,True,1),
        (16,1,80,512,4,16,"expert_slice",True,True,True,1),
        (2,4,80,512,8,8,"plain",True,False,True,1),
        (1,1,80,512,8,10,"rank2",False,True,False,1),
        (4,1,80,512,8,40,"input_slice",False,True,False,3),
        (32,1,80,512,4,32,"transpose",False,True,True,1),
    ]
    for bits,group in ((4,32),(4,64),(8,32),(8,64)):
        for case,(xb,m,n,k,e,r,layout,left,right,sorted_indices,rank) in enumerate(cases):
            extra = 2 if layout == "expert_slice" else 1
            wp=k*bits//32;sp=k//group
            words=(np.arange(e*n*extra*wp,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
            w=mx.array(words).reshape(e,n*extra,wp)
            s=(pattern(e*n*extra*sp,33)/64).reshape(e,n*extra,sp)
            b=(pattern(e*n*extra*sp,34)/64).reshape(e,n*extra,sp)
            if layout=="expert_slice":w=w[:,n:];s=s[:,n:];b=b[:,n:]
            x=pattern(xb*m*k*(2 if layout=="input_slice" else 1),32)
            if layout=="transpose":
                x=x.reshape(xb,k,m).transpose(0,2,1)
                w=w.reshape(e,wp,n).transpose(0,2,1)
                s=s.reshape(e,sp,n).transpose(0,2,1)
                b=b.reshape(e,sp,n).transpose(0,2,1)
            elif layout=="input_slice":x=x.reshape(xb,m,k*2)[...,::2]
            else:x=x.reshape(xb,m,k)
            if rank==2:x=x.reshape(1,xb,1,m,k)
            if rank==3:x=x.reshape(2,xb//2,1,m,k)
            if layout=="rank2":x=x.reshape(m,k)
            rhs=(np.arange(r,dtype=np.uint32)*5+2)%e
            if sorted_indices:rhs=np.sort(rhs)
            if rank>1:
                raw=np.zeros((xb,r//xb+3),np.int32);raw[:,1:1+r//xb]=rhs.reshape(xb,-1)
                rhs=mx.array(raw)[:,1:1+r//xb]
                rhs=rhs.reshape((1,xb,r//xb) if rank==2 else (2,xb//2,r//xb))
            else:rhs=mx.array(rhs,mx.int32)
            lhs=mx.array((np.arange(r)*2+1)%xb,mx.int64) if left else None
            out=mx.gather_qmm(x,w,s,b,lhs_indices=lhs,rhs_indices=rhs if right else None,
                transpose=True,group_size=group,bits=bits,sorted_indices=sorted_indices)
            emit(f"gather_array_{case}_b{bits}_g{group}",[out],case=case,xb=xb,m=m,n=n,k=k,e=e,r=r,
                 layout=layout,left=left,right=right,sorted=sorted_indices,rank=rank,bits=bits,group=group,shape=out.shape)


def emit(name, outputs, **metadata):
    mx.eval(*outputs)
    # Some original-runtime empty Matmul results have a null data pointer;
    # numpy.asarray crashes although mx.eval succeeded. No bytes exist to read.
    print(json.dumps({"kernel": name, "sha256": [hashlib.sha256(
        b"" if value.size == 0 else np.asarray(value.astype(mx.float32)).astype('<f4').tobytes()).hexdigest()
        for value in outputs], **metadata}, separators=(",", ":")), flush=True)


def quantized_linear_fixture(n, k, bits, group, salt):
    import mlx.nn as nn
    layer = nn.QuantizedLinear(64, 1, bias=False, group_size=group, bits=bits)
    words = (np.arange(n*k*bits//32,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
    layer.weight = mx.array(words).reshape(n,k*bits//32)
    layer.scales = (pattern(n*k//group,salt)/64).reshape(n,k//group)
    layer.biases = (pattern(n*k//group,salt+1)/64).reshape(n,k//group)
    return layer


def hyper_graph_fixtures():
    """Original quantizer + HC pack/primitives, with raw uint32 pack hashes."""
    from types import SimpleNamespace
    from mtplx.kernels.hyper_connection_v3 import prepare_v3_pack, device_supports_hyper_v3
    assert mx.__version__ == "0.32.2"
    assert device_supports_hyper_v3()
    def packed(name, outputs, **metadata):
        mx.eval(*outputs)
        hashes = [hashlib.sha256(np.asarray(value if value.dtype == mx.uint32 else value.astype(mx.float32))
                  .astype('<u4' if value.dtype == mx.uint32 else '<f4').tobytes()).hexdigest() for value in outputs]
        print(json.dumps(dict(kernel=name,sha256=hashes,**metadata),separators=(',',':')),flush=True)
    for dtype in (mx.bfloat16,mx.float16,mx.float32):
        for group in (32,64,128):
            for bits in (2,3,4,5,6,8):
                x = pattern(3*256,3).astype(dtype).reshape(3,256)
                packed(f"quantize_{str(dtype).split('.')[-1]}_g{group}_b{bits}_plain",mx.quantize(x,group_size=group,bits=bits),
                       dtype=str(dtype).split('.')[-1],group=group,bits=bits,layout="plain")
    for layout,x in (
        ("slice",pattern(3*512,3).reshape(3,512)[:,::2]),
        ("transpose",pattern(256*3,3).reshape(256,3).T),
        ("zero",mx.zeros((3,256),mx.bfloat16)),
        ("positive",mx.full((3,256),1,mx.bfloat16)),
        ("negative",mx.full((3,256),-1,mx.bfloat16)),
    ):
        packed(f"quantize_bfloat16_g64_b8_{layout}",mx.quantize(x,group_size=64,bits=8),dtype="bfloat16",group=64,bits=8,layout=layout)
    down,up,inject = pattern(320*10240,3).reshape(320,10240),pattern(10240*320,4).reshape(10240,320),pattern(4*10240,5).reshape(4,10240)
    module=SimpleNamespace(input_mix_weight_down=SimpleNamespace(weight=down),input_mix_weight_up=SimpleNamespace(weight=up),block_inject_weight=SimpleNamespace(weight=inject))
    pack=prepare_v3_pack(module)
    packed("hyper_v3_original_pack",pack)
    emit("hyper_v3_original_packed_read",fused_hyper_read_v3(pattern(10240,1),pattern(10240,2),pack))


def dequantize_fixtures():
    """Original affine dequantizer with all supported table formats and views."""
    assert mx.__version__ == "0.32.2"
    for dtype in (mx.bfloat16, mx.float16, mx.float32):
        for group in (32,64,128):
            for bits in (2,3,4,5,6,8):
                for layout in ("plain","transpose","slice","offset","broadcast","batch"):
                    width=256*bits//32
                    def value(cols, packed=False, salt=3):
                        shape = (cols,3) if layout=="transpose" else (3,cols*2) if layout=="slice" else (4,cols) if layout=="offset" else (1,cols) if layout=="broadcast" else (2,3,cols) if layout=="batch" else (3,cols)
                        count=int(np.prod(shape))
                        a = mx.array((np.arange(count,dtype=np.uint64)*2654435761+12345).astype(np.uint32)) if packed else (pattern(count,salt)/64).astype(dtype)
                        a=a.reshape(shape)
                        if layout=="transpose": a=a.T
                        elif layout=="slice": a=a[:,::2]
                        elif layout=="offset": a=a[1:]
                        elif layout=="broadcast": a=mx.broadcast_to(a,(3,cols))
                        return a
                    w,s,b=value(width,True),value(256//group),value(256//group,salt=4)
                    emit(f"dequantize_{str(dtype).split('.')[-1]}_g{group}_b{bits}_{layout}",
                         [mx.dequantize(w,s,b,group_size=group,bits=bits)],
                         dtype=str(dtype).split('.')[-1],group=group,bits=bits,layout=layout)


def ngram_resident_fixtures():
    """Actual NGramTable._lazy_gather, including general Gather view dispatch."""
    from mtplx.models.qwen4_exp import NGramTable
    for dtype in (mx.bfloat16,mx.float16):
        for bits,width in ((0,8),(0,9),(0,160),(4,160)):
            for layout in ("plain","transpose","offset"):
                for ids_layout in ("vector","scalar","matrix","head3","head4","slice","broadcast"):
                    if bits and ids_layout=="scalar": continue  # Original dequantize requires rank>=2.
                    def part(cols,packed=False,salt=3):
                        shape=(cols,17) if layout=="transpose" else (18,cols) if layout=="offset" else (17,cols)
                        count=int(np.prod(shape))
                        a=mx.array((np.arange(count,dtype=np.uint64)*2654435761+12345).astype(np.uint32)) if packed else (pattern(count,salt)/64).astype(dtype)
                        a=a.reshape(shape)
                        return a.T if layout=="transpose" else a[1:] if layout=="offset" else a
                    table=NGramTable(17,width,sidecar=True)
                    table._lazy_parts=(part(width*bits//32,True),part(width//32),part(width//32,salt=4)) if bits else (part(width),None,None)
                    table._lazy_bits=bits;table._lazy_group=32
                    shapes={"vector":(7,),"scalar":(),"matrix":(2,7),"head3":(2,4,16),"head4":(2,2,4,16),"slice":(2,9,16),"broadcast":(1,1,16)}
                    shape=shapes[ids_layout];count=int(np.prod(shape))
                    ids=mx.array(((np.arange(count,dtype=np.int64)*7)%34)-17).reshape(shape)
                    if ids_layout=="slice": ids=ids[:,1::2,:]
                    if ids_layout=="broadcast": ids=mx.broadcast_to(ids,(2,4,16))
                    emit(f"ngram_resident_{str(dtype).split('.')[-1]}_b{bits}_w{width}_{layout}_{ids_layout}",[table._lazy_gather(ids)],
                         dtype=str(dtype).split('.')[-1],bits=bits,width=width,layout=layout,ids_layout=ids_layout)


def hyper_forward_fixtures():
    """Actual GatedResidual, including selection and persistent module pack."""
    from mtplx.models.qwen4_exp import GatedResidual, TextArgs
    from mtplx.kernels import hyper_connection as hc, hyper_connection_v3 as v3
    calls=[]
    original_read,original_v3=hc.fused_hyper_read,v3.fused_hyper_read_v3
    def read(*args,**kwargs):
        calls.append("fused")
        return original_read(*args,**kwargs)
    def read_v3(*args,**kwargs):
        calls.append("v3")
        return original_v3(*args,**kwargs)
    hc.fused_hyper_read,v3.fused_hyper_read_v3=read,read_v3
    try:
        for layout in ("dense","quantized"):
            for combine in (False,True):
                layer=GatedResidual(TextArgs(),use_combine=combine)
                layer.hc_norm.weight=pattern(10240,2)
                names=[("input_mix_weight_down",320,10240,3),("input_mix_weight_up",10240,320,4)]
                if combine: names.append(("block_inject_weight",4,10240,5))
                for name,n,k,salt in names:
                    if layout=="dense": getattr(layer,name).weight=pattern(n*k,salt).reshape(n,k)
                    else: setattr(layer,name,quantized_linear_fixture(n,k,4,64,salt))
                for rows in (1,2,4,8,15,16,32,2048):
                    for mode in range(4):
                        os.environ["MTPLX_FUSED_HC"]=str(mode & 1)
                        os.environ["MTPLX_FUSED_HC_V3"]=str(bool(mode & 2)).lower()
                        x=pattern(rows*10240,1).reshape(1,rows,10240)
                        calls.clear()
                        output=layer(x)
                        selected=calls[-1] if calls else "eager"
                        if combine:
                            mixed,hyper,inject=output
                            assert hyper is x
                            output=[mixed,inject]
                        else: output=[output]
                        emit(f"hyper_forward_{layout}_combine{int(combine)}_r{rows}_mode{mode}",output,
                             layout=layout,combine=combine,rows=rows,mode=mode,selected=selected,pack_exists=getattr(layer,"_v3_pack",None) is not None)
    finally:
        hc.fused_hyper_read,v3.fused_hyper_read_v3=original_read,original_v3


def ngram_rows_fixtures():
    """Actual NGramEmbedding CPU stage hash; raw I64, EOS and continuation."""
    from mtplx.models.qwen4_exp import NGramEmbedding, TextArgs
    layer=NGramEmbedding(TextArgs(ngram_sidecar=True,eos_token_id=248044),0)
    mult,sizes,offsets=layer._np_consts()
    print(json.dumps(dict(contract=True,multipliers=mult.tolist(),sizes=sizes.tolist(),
                          offsets=offsets.tolist(),eos=layer.eos_id)),flush=True)
    def ids(n,salt):
        palette=np.array([0,1,42,layer.eos_id,layer.eos_id,248000,-1,2**63-1,-2**63,43],dtype=np.int64)
        return palette[(np.arange(n)*7+salt)%len(palette)]
    for batch in (1,2):
        for rows in (0,1,4,2048):
            for prev_len in (0,1,2,5):
                prev=ids(batch*prev_len,3).reshape(batch,prev_len)
                for step,tokens in enumerate((rows,1)):
                    values=ids(batch*tokens,step).reshape(batch,tokens)
                    result,history=layer._rows_np(values,prev)
                    print(json.dumps(dict(kernel=f"ngram_rows_b{batch}_r{rows}_p{prev_len}_step{step}",
                        batch=batch,rows=rows,prev_len=prev_len,step=step,
                        result_shape=list(result.shape),history_shape=list(history.shape),
                        sha256=[hashlib.sha256(np.asarray(v).astype('<i8').tobytes()).hexdigest()
                                for v in (result,history)]),separators=(',',':')),flush=True)
                    prev=history


def ngram_gpu_fixtures():
    """Actual _graph_path through row IDs; table returns IDs, not embeddings."""
    import mlx.nn as nn
    from mlx_lm.models.cache import ArraysCache
    from mtplx.models.qwen4_exp import NGramEmbedding, TextArgs
    class RowsOnly(nn.Module):
        def __call__(self,ids): return ids
    layer=NGramEmbedding(TextArgs(ngram_sidecar=True,eos_token_id=248044),0)
    layer.ngram_embedding=RowsOnly()
    palette=np.array([0,1,42,layer.eos_id,layer.eos_id,248000,-1,2**63-1,-2**63,43],dtype=np.int64)
    def ids(n,salt): return palette[(np.arange(n)*7+salt)%len(palette)]
    for batch in (1,2):
        for rows in (1,4,2048):
            for prev_len in (-1,0,1,2,5):
                cache=ArraysCache(4)
                if prev_len>=0: cache[3]=mx.array(ids(batch*prev_len,3).reshape(batch,prev_len))
                for step,tokens in enumerate((rows,1)):
                    values=mx.array(ids(batch*tokens,step).reshape(batch,tokens))
                    result=layer._graph_path(values,cache,3)
                    mx.eval(result,cache[3])
                    print(json.dumps(dict(kernel=f"ngram_gpu_b{batch}_r{rows}_p{prev_len}_step{step}",
                        batch=batch,rows=rows,prev_len=prev_len,step=step,
                        result_shape=list(result.shape),history_shape=list(cache[3].shape),
                        sha256=[hashlib.sha256(np.asarray(v).astype('<i8').tobytes()).hexdigest()
                                for v in (result,cache[3])]),separators=(',',':')),flush=True)


def ngram_sidecar_fixtures():
    """Actual file-backed _SidecarGather, including persistent LRU/bypass state."""
    import tempfile
    from mtplx.models.qwen4_exp import _SidecarGather
    cases=([5,1,5,3],[3,5,2,1],[6,2,9,6],[],list(range(4095,-1,-1)),list(range(4096,-1,-1)),[-1,5002,-2],[1,3,5])
    for dtype in ("BF16","F16"):
        for bits in (0,4):
            with tempfile.TemporaryDirectory(prefix="mtplx-sidecar-") as temp:
                path=Path(temp)/"rows.bin"
                entries={};offset=4096
                names=("weight",) if bits==0 else ("weight","scales","biases")
                with path.open("wb") as f:
                    for part,name in enumerate(names):
                        packed=bits!=0 and name=="weight"
                        columns=20 if packed else 5 if bits else 160
                        dt="U32" if packed else dtype
                        count=5003*columns
                        values=(np.arange(count,dtype=np.uint64)*2654435761+12345).astype('<u4') if packed else (0x3c00+(np.arange(count,dtype=np.uint64)*7+part*13)%512).astype('<u2')
                        f.seek(offset);f.write(values.tobytes())
                        entries[name]=({"dtype":dt,"shape":[5003,columns],"data_offsets":[offset,offset+values.nbytes]},0)
                        offset+=values.nbytes+34
                for prefetch in (False,True):
                    os.environ["MTPLX_NGRAM_PREFETCH"]="1" if prefetch else "0"
                    for capacity in (0,3,4096):
                        gather=_SidecarGather(path,entries,bits,32);gather._hot_cap_rows=capacity
                        try:
                            for step,ids in enumerate(cases):
                                out=gather(mx.array(ids,mx.int64),160);mx.eval(out)
                                keys=np.array(list(gather._hot),np.int64)
                                print(json.dumps(dict(kernel=f"ngram_sidecar_{dtype}_b{bits}_p{int(prefetch)}_c{capacity}_s{step}",dtype=dtype,bits=bits,prefetch=prefetch,capacity=capacity,step=step,
                                    sha256=hashlib.sha256(np.asarray(out.astype(mx.float32)).astype('<f4').tobytes()).hexdigest(),
                                    order_sha256=hashlib.sha256(keys.astype('<i8').tobytes()).hexdigest(),entries=len(keys),hits=gather.hot_hits,misses=gather.hot_misses,prefetch_batches=gather.prefetch_batches),separators=(',',':')),flush=True)
                        finally:
                            if gather._pool is not None: gather._pool.shutdown()
                            os.close(gather._fd)


def ngram_stage_fixtures():
    """Original stage/__call__ lifecycle with real file-backed small tables."""
    import tempfile,contextlib,io
    from mlx_lm.models.cache import ArraysCache
    from mtplx.models.qwen4_exp import NGramEmbedding,TextArgs,_SidecarGather
    scenarios=("normal","verify","graph","disabled","stale","different","mismatch","skip_pending","error","no_cache","seeded","overwrite","graph_error","missing_skip","skip_policy")
    for bits in (0,4):
        with tempfile.TemporaryDirectory(prefix="mtplx-stage-") as temp:
            path=Path(temp)/"rows.bin";entries={};offset=4096
            with path.open("wb") as f:
                for part,name in enumerate(("weight",) if bits==0 else ("weight","scales","biases")):
                    packed=bits and name=="weight";cols=20 if packed else 5 if bits else 160
                    values=(np.arange(17*cols,dtype=np.uint64)*2654435761+12345).astype('<u4') if packed else (0x3c00+(np.arange(17*cols,dtype=np.uint64)*7+part*13)%512).astype('<u2')
                    f.seek(offset);f.write(values.tobytes());entries[name]=({"dtype":"U32" if packed else "BF16","shape":[17,cols],"data_offsets":[offset,offset+values.nbytes]},0);offset+=values.nbytes+34
            for scenario in scenarios:
                for batch,rows in ((1,4),(2,4),*(([(1,2048),(2,2048)]) if scenario in ("normal","verify") else [])):
                    layer=NGramEmbedding(TextArgs(ngram_sidecar=True,eos_token_id=248044),0)
                    layer.ngram_heads_vocab_sizes=mx.full((16,),17,mx.int64);layer.ngram_heads_offsets=mx.zeros((16,),mx.int64)
                    os.environ["MTPLX_NGRAM_PREFETCH"]="0"
                    sidecar=_SidecarGather(path,entries,bits,32);layer.ngram_embedding._sidecar=sidecar
                    cache=None if scenario=="no_cache" else ArraysCache(4)
                    if scenario=="seeded": cache[3]=mx.arange(batch*5,dtype=mx.int64).reshape(batch,5)
                    os.environ["MTPLX_NGRAM_STAGE"]="0" if scenario in ("graph","graph_error") else "1"
                    os.environ["MTPLX_NGRAM_STAGE_VERIFY"]="1" if scenario in ("verify","mismatch") else "0"
                    layer._stage_disabled=scenario=="disabled"
                    try:
                        for step,tokens in enumerate((rows,1)):
                            ids=mx.array((np.arange(batch*tokens)*7+step)%51,mx.int32).reshape(batch,tokens)
                            alternate=ids+1
                            with contextlib.redirect_stdout(io.StringIO()):
                                layer.stage(ids,cache,3)
                                if step==0:
                                    if scenario=="skip_pending":
                                        layer._stage_disabled=True;layer.stage(alternate,cache,3);layer._stage_disabled=False
                                    if scenario=="missing_skip":
                                        layer.ngram_embedding._sidecar=None;layer.stage(alternate,cache,3);layer.ngram_embedding._sidecar=sidecar
                                    if scenario=="skip_policy":
                                        os.environ["MTPLX_NGRAM_STAGE"]="0";layer.stage(alternate,cache,3);os.environ["MTPLX_NGRAM_STAGE"]="1"
                                    if scenario=="error": layer.stage(ids.reshape(-1),cache,3);layer.stage(ids.reshape(-1),cache,3)
                                    if scenario=="overwrite": layer.stage(alternate,cache,3)
                            prehistory=None if cache is None or cache[3] is None else hashlib.sha256(np.asarray(cache[3]).astype('<i8').tobytes()).hexdigest()
                            called=ids[:,:1] if step==0 and scenario=="stale" else alternate if step==0 and scenario in ("different","mismatch","skip_pending","overwrite","missing_skip","skip_policy") else ids
                            if step==0 and scenario=="graph_error": layer.ngram_embedding._sidecar=None
                            error=False;out=None
                            try: out=layer(called,cache,3);mx.eval(out)
                            except RuntimeError as exc:
                                assert "staged/graph mismatch" in str(exc) or "sidecar was never attached" in str(exc);error=True
                            layer.ngram_embedding._sidecar=sidecar
                            history=None if cache is None or cache[3] is None else hashlib.sha256(np.asarray(cache[3]).astype('<i8').tobytes()).hexdigest()
                            print(json.dumps(dict(kernel=f"ngram_stage_b{bits}_{scenario}_B{batch}_r{rows}_s{step}",bits=bits,scenario=scenario,batch=batch,rows=rows,step=step,error=error,
                                sha256=None if out is None else hashlib.sha256(np.asarray(out.astype(mx.float32)).astype('<f4').tobytes()).hexdigest(),history=history,prehistory=prehistory,
                                consumed=getattr(layer,"_stage_consumed",0),bypassed=getattr(layer,"_stage_bypassed",0),graph_calls=getattr(layer,"_graph_calls",0),warned=getattr(layer,"_stage_warned",False),pending=getattr(layer,"_staged",None) is not None),separators=(',',':')),flush=True)
                    finally: os.close(sidecar._fd)


def ngram_stage_verification_fixtures():
    """Original QA allclose special values and all-reduction size boundaries."""
    for size in (0,1,2,4096,4097,1<<20):
        for mode in range(6):
            a=np.ones(size,np.float32);b=a.copy()
            if size:
                if mode==1: b[-1]=2
                if mode==2: a[-1]=b[-1]=np.inf
                if mode==3: a[-1]=b[-1]=-np.inf
                if mode==4: a[-1]=np.inf;b[-1]=-np.inf
                if mode==5: a[-1]=b[-1]=np.nan
            print(json.dumps(dict(kernel=f"stage_allclose_{size}_{mode}",value=bool(mx.allclose(mx.array(a),mx.array(b)))),separators=(',',':')),flush=True)
    for size in (1<<26,(1<<26)+1):
        print(json.dumps(dict(kernel=f"stage_all_{size}",value=bool(mx.all(mx.array(np.ones(size,np.bool_))))),separators=(',',':')),flush=True)


def ngram_embedding_fixtures():
    """Entire original resident _graph_path with a small table, actual head width."""
    from mlx_lm.models.cache import ArraysCache
    from mtplx.models.qwen4_exp import NGramEmbedding,TextArgs
    layer=NGramEmbedding(TextArgs(ngram_sidecar=True,eos_token_id=248044),0)
    layer.ngram_heads_vocab_sizes=mx.full((16,),17,mx.int64)
    layer.ngram_heads_offsets=mx.zeros((16,),mx.int64)
    table=layer.ngram_embedding
    table.prefer_lazy=True;table._lazy_group=32
    for bits in (0,4):
        table._lazy_bits=bits
        if bits:
            words=(np.arange(17*20,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
            table._lazy_parts=(mx.array(words).reshape(17,20),(pattern(17*5,3)/64).reshape(17,5),(pattern(17*5,4)/64).reshape(17,5))
        else: table._lazy_parts=((pattern(17*160,3)/64).reshape(17,160),None,None)
        for batch in (1,2):
            for rows in (1,4,2048):
                cache=ArraysCache(4)
                for step,tokens in enumerate((rows,1)):
                    values=mx.array(((np.arange(batch*tokens)*7+step)%51),mx.int64).reshape(batch,tokens)
                    result=layer._graph_path(values,cache,3)
                    mx.eval(result,cache[3])
                    print(json.dumps(dict(kernel=f"ngram_embedding_b{bits}_batch{batch}_r{rows}_step{step}",bits=bits,batch=batch,rows=rows,step=step,
                        sha256=[hashlib.sha256(np.asarray(result.astype(mx.float32)).astype('<f4').tobytes()).hexdigest(),hashlib.sha256(np.asarray(cache[3]).astype('<i8').tobytes()).hexdigest()]),separators=(',',':')),flush=True)


def ple_connected_fixtures():
    """Unmodified PLELayer including real lookup, stage and both cache slots."""
    import tempfile, contextlib, io
    from mlx_lm.models.cache import ArraysCache
    from mtplx.models.qwen4_exp import PLELayer, TextArgs
    assert mx.__version__ == "0.32.2"
    for bits in (0,4):
        layer=PLELayer(TextArgs(ngram_sidecar=True,eos_token_id=248044),0)
        for name,n,salt in (("key_proj",10240,71),("value_proj",2560,73)):
            if bits==0: getattr(layer,name).weight=pattern(n*2560,salt).reshape(n,2560)
            else: setattr(layer,name,quantized_linear_fixture(n,2560,4,64,salt))
        for name,salt in (("norm_key",75),("norm_query",76),("norm_conv",77)):
            getattr(layer,name).weight=pattern(10240,salt)
        layer.conv_weight=pattern(10240*4,78).reshape(10240,4,1)
        embedding=layer.ple_embedding
        embedding.ngram_heads_vocab_sizes=mx.full((16,),17,mx.int64)
        embedding.ngram_heads_offsets=mx.zeros((16,),mx.int64)
        with tempfile.TemporaryDirectory(prefix="mtplx-ple-") as temp:
            path=Path(temp)/"rows.safetensors";header={"__metadata__":{"ngram_bits":str(bits),"ngram_group_size":"32"}};offset=4096
            with path.open("wb") as f:
                for part,name in enumerate(("weight",) if bits==0 else ("weight","scales","biases")):
                    packed=bits and name=="weight";cols=20 if packed else 5 if bits else 160
                    values=(np.arange(17*cols,dtype=np.uint64)*2654435761+12345).astype('<u4') if packed else (0x3c00+(np.arange(17*cols,dtype=np.uint64)*7+part*13)%512).astype('<u2')
                    f.seek(offset);f.write(values.tobytes());header[f"ngram.{name}"]={"dtype":"U32" if packed else "BF16","shape":[17,cols],"data_offsets":[offset-4096,offset-4096+values.nbytes]};offset+=values.nbytes+34
                encoded=json.dumps(header).encode();assert len(encoded)<=4088
                f.seek(0);f.write((4088).to_bytes(8,"little"));f.write(encoded.ljust(4088,b' '))
            for mode in ("stage","verify","graph","resident"):
                os.environ["MTPLX_NGRAM_STAGE"]="0" if mode=="graph" else "1"
                os.environ["MTPLX_NGRAM_STAGE_VERIFY"]="1" if mode=="verify" else "0"
                os.environ["MTPLX_NGRAM_PREFETCH"]="0"
                table=embedding.ngram_embedding
                table.prefer_lazy=mode=="resident"
                embedding._stage_disabled=mode=="resident"
                for batch,rows in ((1,7),(1,2048),(2,7)):
                    for state in ("none","empty","seeded"):
                        table.attach_sidecar(path);sidecar=table._sidecar
                        if mode=="resident":
                            log=io.StringIO()
                            with contextlib.redirect_stdout(log): bound=table.attach_resident(path)
                            assert bound, log.getvalue()
                            assert not table.prefer_lazy
                            table.prefer_lazy=True
                        cache=None if state=="none" else ArraysCache(4)
                        if state=="seeded":
                            cache[2]=pattern(batch*12*10240,79).reshape(batch,12,10240)
                            cache[3]=mx.arange(batch*5,dtype=mx.int64).reshape(batch,5)
                        try:
                            for step,tokens in enumerate((rows,1)):
                                values=((np.arange(batch*tokens)*7+step)%51).astype(np.int32)
                                values[min(2,len(values)-1)]=248044
                                ids=mx.array(values).reshape(batch,tokens)
                                hidden=pattern(batch*tokens*10240,80+step).reshape(batch,tokens,10240)
                                embedding.stage(ids,cache,3)
                                output=layer(hidden,ids,cache)
                                mx.eval(output,*([] if cache is None else [cache[2],cache[3]]))
                                history=None if cache is None else hashlib.sha256(np.asarray(cache[3]).astype('<i8').tobytes()).hexdigest()
                                emit(f"ple_connected_b{bits}_{mode}_B{batch}_r{rows}_{state}_s{step}",
                                     [output] if cache is None else [output,cache[2]],
                                     bits=bits,mode=mode,batch=batch,rows=rows,state=state,step=step,history=history)
                        finally: os.close(sidecar._fd)


def batched_sampler_fixtures():
    """Both actual batch routes, their whole-batch fallback and public methods."""
    from mtplx import fast_sampling as fs
    from mtplx.sampling import SamplerConfig
    assert mx.__version__ == "0.32.2" and np.__version__ == "2.4.4"
    def digest(values, dtype):
        return hashlib.sha256(np.asarray(values, dtype=dtype).tobytes()).hexdigest()
    def bits(value): return int(np.asarray(value, dtype='<f8').view('<u8'))
    def receipt(batch):
        if batch is None: return dict(none=True)
        rng = np.random.default_rng(43210)
        samples = []
        for i in range(32):
            try: samples.append(batch.sample(i % len(batch.token_ids), rng))
            except ValueError: samples.append(None)
        converted = []
        for row in range(len(batch.token_ids)):
            try:
                p = batch.to_distribution(row)
                converted.append(dict(ids=digest(p.token_ids, '<i8'), probs=digest(p.probs, '<f8')))
            except ValueError: converted.append(dict(error=True))
        lookups = []
        for row in (0, -1, len(batch.token_ids), -len(batch.token_ids)-1):
            for token in (-1, 0, 3, batch.vocab_size-1, batch.vocab_size):
                try: lookups.append(bits(batch.probability(row, token)))
                except IndexError: lookups.append(None)
        return dict(ids=digest(batch.token_ids, '<i8'), probs=digest(batch.probs, '<f8'),
            shape=list(batch.token_ids.shape), samples=samples, next_rng=bits(rng.random()),
            converted=converted, lookups=lookups, vocab=batch.vocab_size)
    original_eval, original_host = mx.eval, fs._host_sparse_distribution
    trace = []
    host_calls = []
    def evaluate(*args, **kwargs):
        trace.append([list(a.shape) for a in args])
        return original_eval(*args, **kwargs)
    def host(*args, **kwargs):
        host_calls.append(1)
        return original_host(*args, **kwargs)
    mx.eval, fs._host_sparse_distribution = evaluate, host
    try:
        with np.errstate(all="ignore"):
            for vocab in (17, 4097, 248320):
                for dtype in (mx.bfloat16, mx.float16, mx.float32):
                    for layout in ("plain", "strided", "reverse", "broadcast"):
                        for mode in ("pattern", "equal", "mixed"):
                            for temperature, top_p, k in ((0.6, 0.95, 20), (0.7, 1.0, 20), (1e-7, 0.01, 129)):
                                for route in ("serial", "bound"):
                                    cols = 2*vocab if layout == "strided" else vocab
                                    count = cols if layout == "broadcast" else 4*cols
                                    x = (((np.arange(count,dtype=np.int64)*17+96*13)%257)-128).astype(np.float32)/np.float32(131)
                                    if mode == "equal": x[:] = 1
                                    if mode == "mixed":
                                        x[-cols:] = np.nan
                                    row = mx.array(x).astype(dtype) + mx.array(0, dtype=dtype)
                                    row = row.reshape(-1, cols)
                                    if layout == "strided": row = row[:, 1::2]
                                    elif layout == "reverse": row = row[:, ::-1]
                                    elif layout == "broadcast": row = mx.broadcast_to(row, (4, vocab))
                                    config = SamplerConfig(temperature=temperature, top_p=top_p, top_k=k)
                                    trace.clear(); host_calls.clear()
                                    try:
                                        batch = (fs.bind_batched_top_k_distributions(config, vocab_size=vocab)(row)
                                            if route == "bound" else fs.batched_sparse_distributions_from_mlx_logits(row, config))
                                        result = receipt(batch)
                                    except (ValueError, FloatingPointError): result = dict(error=True)
                                    print(json.dumps(dict(vocab=vocab, dtype=str(dtype).split('.')[-1], layout=layout, mode=mode,
                                        temperature=temperature, top_p=top_p, k=k, route=route, result=result,
                                        evals=trace.copy(), host_calls=len(host_calls)), separators=(',',':')), flush=True)
            # Constructor preserves negatives if the total remains positive;
            # sample rejects the unnormalised positive subset without an RNG draw.
            for name, ids, probs in (
                ("padded", [[3,0,-1,7],[0,3,3,-1]], [[1,3,0,2],[0,1,2,0]]),
                ("negative", [[0,1,2],[0,1,2]], [[2,-1,1],[1,0,0]]),
                ("zero", [[0,1]], [[0,0]]),
                ("nonfinite", [[0,1]], [[np.inf,1]]),
                ("mismatch", [[0,1]], [[1,2,3]])):
                try: result=receipt(fs.BatchedSparseDistributions(np.array(ids), np.array(probs), vocab_size=17))
                except ValueError: result=dict(error=True)
                print(json.dumps(dict(constructor=name, result=result), separators=(',',':')), flush=True)
    finally:
        mx.eval, fs._host_sparse_distribution = original_eval, original_host


def mtp_prefill_driver_fixtures():
    """Unchanged committed-history prefill, including abort and history windows."""
    from mtplx import generation as gen
    from mtplx.runtime import MTPLXRuntime
    from mlx_lm.models.cache import ArraysCache
    from types import SimpleNamespace
    import itertools
    for name in list(os.environ):
        if name.startswith("MTPLX_"): del os.environ[name]
    original_tick, original_count = gen._owner_progress_tick, gen._runtime_count
    cases = list(itertools.product((0, 1, 9), (None, 4), (None, 0, 3), (False, True),
        (False, True), ("normal", "cache", "discard"), (0, 3, 5)))
    cases += [(2050, 2048, window, absolute, True, "normal", 0)
        for window in (None, 1, 3, 2049) for absolute in (False, True)]
    try:
        for length, chunk, window, absolute, final_only, route, abort_at in cases:
            os.environ["MTPLX_SUSTAINED_PREFILL"] = "1" if chunk is not None else "0"
            os.environ["MTPLX_PREFILL_CHUNK_SIZE"] = str(chunk or 2048)
            os.environ["MTPLX_TARGET_EMIT_FULL_PREFILL_LOGITS"] = "0" if final_only else "1"
            os.environ["MTPLX_PREFILL_OMLX_EXTERNAL"] = "0" if route == "normal" else "1"
            os.environ["MTPLX_PREFILL_EXTERNAL_EMIT_LOGITS"] = "1" if route == "discard" else "0"
            events, calls, history, chunks, checks = [], [], [], [], 0
            class Model:
                def make_cache(self):
                    events.append("init")
                    return [ArraysCache(1)]
                def make_mtp_cache(self):
                    events.append("mtp_init")
                    return [ArraysCache(1)]
                def __call__(self, ids, cache=None, return_hidden=False, emit_logits=True, logits_keep=None, **kwargs):
                    calls.append((ids, return_hidden, emit_logits, logits_keep or 0))
                    value = ids.astype(mx.float32)
                    cache[0][0] = value + 1
                    logits = value[..., None] if emit_logits else None
                    return (logits, value[..., None]) if return_hidden else logits
                def mtp_update_cache(self, hidden, ids, mtp_cache=None, position_offset=None, **kwargs):
                    history.append((hidden, ids, position_offset))
                    mtp_cache[0][0] = hidden + ids[..., None].astype(mx.float32)
                    return mtp_cache[0][0]
            def abort():
                nonlocal checks
                checks += 1
                events.append("check")
                return abort_at > 0 and checks == abort_at
            def on_chunk(data):
                events.append("callback")
                assert data["chunk_elapsed_s"] >= 0 and data["elapsed_s"] >= 0
                chunks.append(dict(tokens=data["tokens_done"], total=data["tokens_total"], size=data["chunk_size"]))
                if absolute: raise ValueError("UI callback errors are deliberately ignored")
            def record_count(runtime, key, amount=1):
                original_count(runtime, key, amount)
                if key == "prefill_chunks": events.append("chunk")
                if key == "mtp_history_append_calls": events.append("history")
            gen._owner_progress_tick = lambda: events.append("settled")
            gen._runtime_count = record_count
            rt = MTPLXRuntime(Model(), SimpleNamespace(), Path("."), True,
                SimpleNamespace(concat_order="embedding_hidden", hidden_variant="post_norm"))
            def data(a): return dict(shape=list(a.shape), values=np.array(a).reshape(-1).tolist())
            try:
                cache, logits, hidden, mtp_cache, target_s, history_s, base = gen._prefill_committed_mtp_history_streaming(
                    rt, [i % 37 for i in range(length)], history_window_tokens=window,
                    mtp_position_mode="absolute" if absolute else "cache", abort_check=abort, chunk_callback=on_chunk)
                result = dict(logits=data(logits), hidden=data(hidden), base=base)
            except (ValueError, gen.PostcommitAbort) as error:
                result = dict(error=str(error))
            print(json.dumps(dict(length=length, chunk=chunk, window=window, absolute=absolute,
                final_only=gen._final_logits_prefill_enabled(), route=route, abort_at=abort_at, events=events, chunks=chunks,
                calls=[dict(ids=np.array(ids).reshape(-1).tolist(), hidden=h, emit=e, keep=k) for ids,h,e,k in calls],
                history=[dict(hidden=data(h), ids=np.array(ids).reshape(-1).tolist(), position=p) for h,ids,p in history], result=result)), flush=True)
    finally:
        gen._owner_progress_tick, gen._runtime_count = original_tick, original_count


def prefill_driver_fixtures():
    """Original cold _prefill control flow, cache roots and cancellation."""
    from mtplx import generation as gen
    from mtplx.runtime import MTPLXRuntime
    from mlx_lm.models.cache import ArraysCache
    from mtplx.models.qwen4_exp import QSACache
    from types import SimpleNamespace
    assert mx.__version__ == "0.32.2"
    for name in list(os.environ):
        if name.startswith("MTPLX_"): del os.environ[name]
    arrays = [mx.full((1, 1, 256, 1), i, mx.float32) for i in range(18)]
    gdn = ArraysCache(4)
    gdn.left_padding, gdn.lengths = arrays[:2]
    gdn.cache = arrays[2:6]
    gdn._mtplx_verify_ple = (arrays[17], arrays[17])
    gdn._mtplx_verify_rows = tuple(arrays[6:12])
    qsa = QSACache()
    qsa.kv.keys, qsa.kv.values = arrays[12], arrays[3]
    qsa.kv.offset = 3
    qsa.raw_keys, qsa.pooled, qsa.pooled_f32_t = arrays[14:17]
    roots = gen._tree_mx_arrays([gdn, qsa, gdn])
    print(json.dumps(dict(kind="roots", tags=[next(i for i, a in enumerate(arrays) if a is root) for root in roots],
        shapes=[root.shape for root in roots])), flush=True)
    original_tick = gen._owner_progress_tick
    original_count = gen._runtime_count
    try:
        for length in (0, 1, 2, 9, 2050):
            for chunk in ((None, 2048) if length > 9 else (None, 1, 4)):
                for final_only in (False, True):
                    for route in ("normal", "cache", "discard"):
                        for hidden in (False, True):
                            for abort_at in (0, 1, 3):
                                os.environ["MTPLX_SUSTAINED_PREFILL"] = "1" if chunk is not None else "0"
                                os.environ["MTPLX_PREFILL_CHUNK_SIZE"] = str(chunk or 2048)
                                os.environ["MTPLX_TARGET_EMIT_FULL_PREFILL_LOGITS"] = "0" if final_only else "1"
                                os.environ["MTPLX_PREFILL_OMLX_EXTERNAL"] = "0" if route == "normal" else "1"
                                os.environ["MTPLX_PREFILL_EXTERNAL_EMIT_LOGITS"] = "1" if route == "discard" else "0"
                                events, calls, checks = [], [], 0
                                class Model:
                                    def make_cache(self):
                                        events.append("init")
                                        return [ArraysCache(1)]
                                    def __call__(self, ids, cache=None, return_hidden=False, emit_logits=True, logits_keep=None, **kwargs):
                                        calls.append((ids, return_hidden, emit_logits, logits_keep or 0))
                                        value = ids.astype(mx.float32)
                                        cache[0][0] = value + 1
                                        logits = value[..., None] if emit_logits else None
                                        return (logits, value[..., None]) if return_hidden else logits
                                def abort():
                                    nonlocal checks
                                    checks += 1
                                    events.append("check")
                                    return abort_at > 0 and checks == abort_at
                                gen._owner_progress_tick = lambda: events.append("settled")
                                rt = MTPLXRuntime(Model(), SimpleNamespace(), Path("."), hidden, None)
                                def record_count(runtime, key, amount=1):
                                    original_count(runtime, key, amount)
                                    if key == "prefill_chunks": events.append("chunk")
                                gen._runtime_count = record_count
                                try:
                                    cache, logits, h, seconds = gen._prefill(rt, [i % 37 for i in range(length)], return_hidden=hidden, abort_check=abort)
                                    result = dict(logits=np.array(logits).reshape(-1).tolist(), shape=logits.shape,
                                        hidden=None if h is None else dict(values=np.array(h).reshape(-1).tolist(), shape=h.shape))
                                except (ValueError, gen.PostcommitAbort) as error:
                                    result = dict(error=str(error))
                                print(json.dumps(dict(kind="prefill", length=length, chunk=chunk,
                                    final_only=gen._final_logits_prefill_enabled(), route=route, hidden=hidden, abort_at=abort_at,
                                    events=events, calls=[dict(ids=np.array(ids).reshape(-1).tolist(), hidden=h, emit=e, keep=k) for ids,h,e,k in calls],
                                    chunks=rt.diagnostic_counters.get("prefill_chunks", 0), result=result)), flush=True)
    finally:
        gen._owner_progress_tick = original_tick
        gen._runtime_count = original_count


def ar_driver_fixtures():
    """Unchanged generate_ar with a tiny GPU transition-table model, no weights.

    Isolates orchestration (not model or throughput parity); actual prefill,
    samplers, pipeline and final-state capture are all MTPLX implementations.
    """
    from mtplx.generation import generate_ar
    from mtplx.runtime import MTPLXRuntime
    from mtplx.sampling import SamplerConfig
    from types import SimpleNamespace
    assert mx.__version__ == "0.32.2"
    for name in list(os.environ):
        if name.startswith("MTPLX_"): del os.environ[name]
    vocab = 37
    values = (((np.arange(vocab*vocab, dtype=np.int64)*17 + 96*13) % 257)-128).astype(np.float32)/131
    table = mx.array(values.reshape(vocab, vocab))
    class Model:
        def __init__(self):
            self.calls, self.modes = [], []
        def make_cache(self): return []
        def set_ar_pipeline_mode(self, enabled):
            self.modes.append(bool(enabled))
            return True
        def __call__(self, ids, cache=None, return_hidden=False, emit_logits=True, logits_keep=None):
            self.calls.append((ids, return_hidden))
            out = table[ids]
            return (out, ids.astype(mx.float32)) if return_hidden else out
    tokenizer = SimpleNamespace(decode=lambda ids, **kwargs: " ".join(map(str, ids)))
    for pipeline in (False, True):
        for asynchronous in (False, True):
            for temperature, penalty in ((0., 0.), (.6, 0.), (.6, .4)):
                for limit in (0, 1, 2, 3, 12):
                    for capture in (False, True):
                        # 13/35 are the first/second sampled pipeline tokens:
                        # exercise immediate EOS and EOS with an in-flight next
                        # graph, not just limits and a non-triggering stop set.
                        for stop_ids in (set(), {0, 5, 36}, {13}, {35}):
                            os.environ["MTPLX_AR_PIPELINE"] = "1" if pipeline else "0"
                            os.environ["MTPLX_ASYNC_AR"] = "1" if asynchronous else "0"
                            model = Model()
                            rt = MTPLXRuntime(model, tokenizer, Path("."), False, None)
                            emitted = []
                            out = generate_ar(rt, [1, 2], max_tokens=limit,
                                sampler=SamplerConfig(temperature=temperature, top_p=.95, top_k=20, presence_penalty=penalty),
                                seed=4294967297, stop_token_ids=stop_ids, token_callback=lambda ids: emitted.extend(ids), capture_final_state=capture)
                            # Decode calls only: real prefill consumed [1] then [2].
                            calls = [{"ids": np.array(ids).reshape(-1).tolist(), "hidden": hidden} for ids, hidden in model.calls[2:]]
                            final = None if out.final_state is None else dict(shape=out.final_state.final_logits.shape,
                                sha256=hashlib.sha256(np.array(out.final_state.final_logits).tobytes()).hexdigest())
                            print(json.dumps(dict(pipeline=pipeline, asynchronous=asynchronous, temperature=temperature,
                                penalty=penalty, limit=limit, capture=capture, stops=sorted(stop_ids), tokens=out.tokens,
                                emitted=emitted, calls=calls, modes=model.modes, final=final,
                                verify_calls=out.stats.verify_calls, finish=out.finish_reason)), flush=True)


def prefix_sampler_fixtures():
    """Actual target-prefix entry, interleaved calls and persistent default RNG."""
    from mtplx.fast_sampling import sample_token_ids_from_mlx_logits
    from mtplx.sampling import SamplerConfig
    assert mx.__version__ == "0.32.2"
    configs = ((0.6, .95, 20), (0., .95, 20), (.6, .95, 0),
        (.6, .95, 33), (.6, 1., 0), (.6, .95, 32), (.6, 1., 33),
        (1e-7, .01, 20), (.6, 1., 1), (.6, 1., 999999), (.6, .95, 20))
    for vocab in (17, 4097, 248320):
        for dtype in (mx.bfloat16, mx.float16, mx.float32):
            for layout in ("vector", "singleton", "prefix", "strided", "reverse", "broadcast"):
                shape = (vocab,) if layout == "vector" else ((1, vocab) if layout == "singleton" else (2, 2, vocab))
                size = int(np.prod(shape)) * (2 if layout == "strided" else 1)
                if layout == "broadcast": size = vocab
                for mode in ("pattern", "equal", "nonfinite"):
                    x = (((np.arange(size, dtype=np.int64)*17 + 96*13) % 257)-128).astype(np.float32)/131
                    if mode == "equal": x[:] = 1
                    elif mode == "nonfinite":
                        x[np.arange(size) % 7 == 0] = np.nan
                        x[np.arange(size) % 7 == 1] = np.inf
                        x[np.arange(size) % 7 == 2] = -np.inf
                    row = mx.array(x).astype(dtype)
                    row = row + mx.array(0, dtype=dtype)
                    if layout == "strided": row = row[1::2]
                    elif layout == "reverse": row = row[::-1]
                    if layout == "broadcast": row = mx.broadcast_to(row, shape)
                    else: row = row.reshape(shape)
                    mx.random.seed(4294967297)
                    outputs = []
                    for step, (temperature, top_p, k) in enumerate(configs):
                        if step == 8: mx.random.seed(17)
                        try:
                            token = sample_token_ids_from_mlx_logits(row,
                                SamplerConfig(temperature=temperature, top_p=top_p, top_k=k))
                            if token is None: result = {"none": True}
                            else:
                                mx.eval(token)
                                result = {"tokens": np.array(token).reshape(-1).tolist(), "shape": token.shape}
                        except (ValueError, RuntimeError) as error:
                            result = {"error": str(error)}
                        state = mx.random.state[0]
                        mx.eval(state)
                        result["key"] = np.array(state).tolist()
                        outputs.append(result)
                    print(json.dumps(dict(vocab=vocab, dtype=str(dtype).split('.')[-1], layout=layout,
                        mode=mode, outputs=outputs)), flush=True)


def resident_sampler_fixtures():
    """Original keyed GPU RNG and _mx_lazy_sample, with chained resident keys."""
    from mtplx.generation import _mx_lazy_sample
    from mtplx.sampling import SamplerConfig
    assert mx.__version__ == "0.32.2"
    def digest(a):
        return hashlib.sha256(np.array(a).tobytes()).hexdigest()
    for seed in (0, 4294967297, 18446744073709551615):
        for layout in ("plain", "strided", "reverse"):
            raw = mx.random.key(seed)
            if layout == "strided":
                raw = mx.stack((raw, raw), axis=1).reshape(-1)[::2]
            elif layout == "reverse":
                raw = raw[::-1]
            for shape in ((0,), (1,), (2,), (3,), (127,), (4097,), (2, 3)):
                bits = mx.random.split(raw, num=int(np.prod(shape)))
                uniform = mx.random.uniform(shape=shape, key=raw)
                key = raw
                keys = []
                for _ in range(4):
                    key, sub = mx.random.split(key)
                    keys.append(sub)
                mx.eval(bits, uniform, key, *keys)
                print(json.dumps(dict(kind="rng", seed=seed, layout=layout,
                    shape=shape, bits=digest(bits), uniform=digest(uniform),
                    keys=[np.array(k).tolist() for k in keys], final_key=np.array(key).tolist())), flush=True)
    for vocab in (17, 4097, 248320):
        for dtype in (mx.bfloat16, mx.float16, mx.float32):
            for layout in ("plain", "strided", "reverse", "broadcast"):
                for mode in ("pattern", "equal", "nonfinite"):
                    size = 2*vocab if layout == "strided" else (1 if layout == "broadcast" else vocab)
                    x = (((np.arange(size, dtype=np.int64)*17 + 96*13) % 257)-128).astype(np.float32)/131
                    if mode == "equal":
                        x[:] = 1
                    elif mode == "nonfinite":
                        for i in range(size):
                            if i % 7 < 3:
                                x[i] = (np.nan, np.inf, -np.inf)[i % 7]
                    row = mx.array(x).astype(dtype)
                    row = row + mx.array(0, dtype=dtype)
                    if layout == "strided": row = row[1::2]
                    elif layout == "reverse": row = row[::-1]
                    elif layout == "broadcast": row = mx.broadcast_to(row, (vocab,))
                    for temperature, top_p, k in ((0.6, 0.95, min(20, vocab-1)), (1.0, 1.0, 2), (1e-7, 0.01, min(129, vocab-1))):
                        config = SamplerConfig(temperature=temperature, top_p=top_p, top_k=k)
                        key = mx.random.key(17)
                        tokens = []
                        for _ in range(4):
                            key, sub = mx.random.split(key)
                            tokens.append(_mx_lazy_sample(row, config, sub))
                        mx.eval(key, *tokens)
                        print(json.dumps(dict(kind="sampler", vocab=vocab, dtype=str(dtype).split('.')[-1],
                            layout=layout, mode=mode, temperature=temperature, top_p=top_p, k=k,
                            tokens=[int(t.item()) for t in tokens], final_key=np.array(key).tolist())), flush=True)


def sampler_entry_fixtures():
    """Actual MTPLX generation entry, three consecutive draws and count updates."""
    from mtplx import generation
    from mtplx.sampling import SamplerConfig, SparseDistribution
    assert mx.__version__ == "0.32.2"
    original_eval, original_argmax = mx.eval, mx.argmax
    trace = []
    def evaluate(*args,**kwargs):
        trace.append("eval")
        return original_eval(*args,**kwargs)
    def argmax(*args,**kwargs):
        trace.append("argmax")
        return original_argmax(*args,**kwargs)
    mx.eval, mx.argmax = evaluate, argmax
    try:
        for vocab in (17,4097,248320):
            for dtype in (mx.bfloat16,mx.float16,mx.float32):
                for layout in ("plain","rank3","strided","reverse","broadcast"):
                    for mode in ("pattern","equal","nonfinite"):
                        size = 2*vocab if layout == "strided" else vocab
                        values = (((np.arange(size,dtype=np.int64)*17+96*13)%257)-128).astype(np.float32)/np.float32(131)
                        if mode == "equal": values[:] = 1
                        if mode == "nonfinite":
                            values[::7]=np.nan; values[1::7]=np.inf; values[2::7]=-np.inf
                        for profile in ("greedy","sparse","dense_top_p","dense_all"):
                            for penalties in (False,True):
                                config = SamplerConfig(temperature=0. if profile=="greedy" else 0.6,
                                    top_k=20 if profile in ("greedy","sparse") else 0,
                                    top_p=1. if profile=="dense_all" else 0.95,
                                    presence_penalty=0.17 if penalties else 0., frequency_penalty=0.13 if penalties else 0.)
                                counts={vocab//2:2}; overlay={1:0.10009} if penalties else {}
                                rng=np.random.default_rng(12345); steps=[]
                                for step in range(3):
                                    x=mx.array(values,dtype=dtype)+mx.array(0.,dtype=dtype)
                                    if layout=="rank3": x=x.reshape(1,1,vocab)
                                    if layout=="strided": x=x[1::2]
                                    if layout=="reverse": x=x[::-1]
                                    if layout=="broadcast": x=mx.broadcast_to(x[:1],(vocab,))
                                    trace.clear()
                                    try:
                                        token,p=generation._sample_from_logits(x,config,rng,token_counts=counts,penalty_overlay=overlay)
                                        result=dict(token=token,kind="none" if p is None else "sparse" if isinstance(p,SparseDistribution) else "dense")
                                        if p is not None:
                                            probs=p.probs if isinstance(p,SparseDistribution) else p
                                            result["probabilities"]=hashlib.sha256(probs.astype('<f8').tobytes()).hexdigest()
                                            if isinstance(p,SparseDistribution): result["ids"]=hashlib.sha256(p.token_ids.astype('<i8').tobytes()).hexdigest()
                                        counts[token]=counts.get(token,0)+1
                                    except (ValueError,IndexError): result=dict(error=True)
                                    result["trace"]=list(trace); steps.append(result)
                                print(json.dumps(dict(vocab=vocab,dtype=str(dtype).split('.')[-1],layout=layout,mode=mode,
                                    profile=profile,penalties=penalties,steps=steps,counts=list(counts.items()),
                                    next_rng=int(np.asarray(rng.random(),dtype='<f8').view('<u8'))),separators=(',',':')),flush=True)
    finally:
        mx.eval, mx.argmax = original_eval, original_argmax


def sampler_penalty_fixtures():
    from mtplx.fast_sampling import apply_penalties_mlx
    from mtplx.sampling import apply_penalties
    assert mx.__version__ == "0.32.2"
    for vocab in (17,4097,248320):
        for dtype in (mx.bfloat16,mx.float16,mx.float32):
            for layout in ("plain","strided","offset","reverse","broadcast"):
                size = 2*vocab if layout == "strided" else vocab+1 if layout == "offset" else vocab
                values = (((np.arange(size,dtype=np.int64)*17+96*13)%257)-128).astype(np.float32)/np.float32(131)
                x = mx.array(values,dtype=dtype)
                if layout == "strided": x = x[1::2]
                if layout == "offset": x = x[1:]
                if layout == "reverse": x = x[::-1]
                if layout == "broadcast": x = mx.broadcast_to(x[:1],(vocab,))
                mx.eval(x)
                raw = np.asarray(x.astype(mx.float32)).astype(np.float64)
                counts = {0:0,vocab//2:3,vocab-1:1}
                overlay = {1:0.00031,vocab//2:0.10009,-1:3.14159}
                profiles = [("noop",counts,0.,0.,{}),("no_counts",{},1.3,0.4,{}),
                    ("counts",counts,0.17,0.13,{}),("overlay",{},0.,0.,overlay),
                    ("both",counts,-0.7,1.3,overlay),("clamped",counts,4.5,-3.7,overlay),
                    ("zero_overlay",{},0.,0.,{vocab//2:0.}),
                    ("single",{vocab//2:2},0.4,0.13,{})]
                for profile,c,p,f,o in profiles:
                    gpu = apply_penalties_mlx(x,c,p,f,o)
                    cpu = apply_penalties(raw,c,p,f,o)
                    mx.eval(gpu)
                    result = np.asarray(gpu.astype(mx.float32)).astype('<f4')
                    print(json.dumps(dict(vocab=vocab,dtype=str(dtype).split('.')[-1],layout=layout,profile=profile,
                        counts=list(c.items()),presence=p,frequency=f,overlay=list(o.items()),
                        gpu_identity=gpu is x,cpu_identity=cpu is raw,
                        gpu=hashlib.sha256(result.tobytes()).hexdigest(),
                        cpu=hashlib.sha256(cpu.astype('<f8').tobytes()).hexdigest()),separators=(',',':')),flush=True)


def sampler_mixed_distribution_fixtures():
    """Original dense/sparse choice, residual, and verify including RNG errors."""
    from mtplx.sampling import SparseDistribution, sample_from_distribution, residual_distribution, verify_one_token
    assert np.__version__ == "2.4.4"
    def bits(value): return int(np.asarray(value,dtype='<f8').view('<u8'))
    def digest(values,dtype): return hashlib.sha256(np.asarray(values,dtype=dtype).tobytes()).hexdigest()
    with np.errstate(all="ignore"):
        for vocab in (1,17,129,4097,248320):
            for mode in ("reverse","same","disjoint","nonfinite","zero"):
                p = ((np.arange(vocab,dtype=np.int64)*17+3)%31).astype(np.float64)/31
                q = p[::-1].copy()
                if mode == "same": q = p.copy()
                if mode == "disjoint":
                    p[1::2] = 0
                    q[::2] = 0
                if mode == "nonfinite":
                    p[0] = np.nan
                    q[-1] = np.inf
                if mode == "zero": p[:] = 0
                for kinds in ("dd","ds","sd","ss"):
                    def make(raw,kind):
                        count = min(20,vocab)
                        ids = np.arange(count,dtype=np.int64)*(vocab-1)//max(1,count-1)
                        return SparseDistribution(ids,raw[ids],vocab) if kind == "s" else raw/raw.sum()
                    target,draft = make(p,kinds[0]),make(q,kinds[1])
                    rng = np.random.default_rng(12345)
                    try:
                        sampled = dict(tokens=[sample_from_distribution(target,rng) for _ in range(32)])
                    except ValueError: sampled = dict(error=True)
                    sampled["next"] = bits(rng.random())
                    try:
                        r = residual_distribution(target,draft)
                        sparse = isinstance(r,SparseDistribution)
                        residual = dict(kind="s" if sparse else "d",borrowed=r is target,
                            probabilities=digest(r.probs if sparse else r,'<f8'))
                        if sparse: residual["ids"] = digest(r.token_ids,'<i8')
                    except ValueError: residual = dict(error=True)
                    rng = np.random.default_rng(43210)
                    decisions=[]
                    for i in range(32):
                        token=(i*7+5)%vocab
                        try:
                            decision=verify_one_token(target,draft,token,rng)
                            decisions.append([decision.accepted,decision.token_id,bits(decision.accept_probability)])
                        except ValueError: decisions.append(None)
                    print(json.dumps(dict(vocab=vocab,mode=mode,kinds=kinds,sampled=sampled,residual=residual,
                        decisions=decisions,next_verify=bits(rng.random())),separators=(',',':')),flush=True)


def sampler_argmax_fixtures():
    """The pinned runtime ArgReduce used by MTPLX's greedy sampler."""
    assert mx.__version__ == "0.32.2"
    for dtype in (mx.bfloat16, mx.float16, mx.float32):
        for cols in (1, 17, 4096, 4097, 248320):
            for layout in ("plain", "swapped", "strided", "offset", "broadcast", "vector", "middle", "reverse"):
                for mode in ("pattern", "nonfinite"):
                    def values(shape):
                        data = (((np.arange(np.prod(shape), dtype=np.int64) * 17 + 96 * 13) % 257) - 128).astype(np.float32) / np.float32(131)
                        if mode == "nonfinite":
                            data[::7] = np.nan
                            data[1::7] = np.inf
                            data[2::7] = -np.inf
                        return mx.array(data.reshape(shape), dtype=dtype)
                    if layout == "offset":
                        x = values((6*cols+1,))[1:].reshape(2,3,cols)
                    elif layout == "strided":
                        x = values((2,3,2*cols))[:,:,1::2]
                    elif layout == "broadcast":
                        x = mx.broadcast_to(values((cols,)),(2,3,cols))
                    elif layout == "vector":
                        x = values((cols,))
                    elif layout == "swapped":
                        x = mx.transpose(values((2,3,cols)),(1,0,2))
                    elif layout == "middle":
                        x = mx.transpose(values((2,3,cols)),(0,2,1))
                    elif layout == "reverse":
                        x = values((2,3,cols))[:,:,::-1]
                    else:
                        x = values((2,3,cols))
                    axis = 1 if layout == "middle" else -1
                    for keepdims in (False,True):
                        out = mx.argmax(x,axis=axis,keepdims=keepdims)
                        mx.eval(out)
                        print(json.dumps(dict(dtype=str(dtype).split('.')[-1],cols=cols,layout=layout,
                            mode=mode,axis=axis,keepdims=keepdims,shape=list(out.shape),
                            sha256=hashlib.sha256(np.asarray(out).astype('<u4').tobytes()).hexdigest()),separators=(',',':')),flush=True)


def sampler_fallback_fixtures():
    from mtplx.fast_sampling import sparse_distributions_from_mlx_logits
    from mtplx.sampling import SamplerConfig, distribution_from_logits
    assert mx.__version__ == "0.32.2"
    for vocab in (17, 4097, 248320):
        for mode in ("unique", "sharp", "equal", "nonfinite", "mixed"):
            values = sampler_serial_values(vocab, 4, mode)
            for top_p in (0.95, 1.0):
                config = SamplerConfig(temperature=0.6, top_p=top_p, top_k=20)
                distributions = sparse_distributions_from_mlx_logits(mx.array(values).reshape(4,vocab),config)
                print(json.dumps(dict(kind="sparse",vocab=vocab,mode=mode,top_p=top_p,
                    ids=[d.token_ids.tolist() for d in distributions],
                    probabilities=[d.probs.astype('<f8').view('<u8').tolist() for d in distributions]),separators=(',',':')),flush=True)
                for top_k in (0,20):
                    for temperature in (0.0,0.6):
                        try:
                            probs = distribution_from_logits(values[:vocab].astype(np.float64),
                                SamplerConfig(temperature=temperature,top_p=top_p,top_k=top_k))
                            result = dict(sha256=hashlib.sha256(probs.astype('<f8').tobytes()).hexdigest())
                        except ValueError:
                            result = dict(error=True)
                        print(json.dumps(dict(kind="dense",vocab=vocab,mode=mode,top_p=top_p,
                            top_k=top_k,temperature=temperature,**result),separators=(',',':')),flush=True)


def sampler_distribution_fixtures():
    """Original normalization/choice/verify over the captured serial support."""
    from mtplx.fast_sampling import _serial_row_distribution
    from mtplx.sampling import SparseDistribution, sample_from_distribution, verify_one_token, residual_distribution
    assert np.__version__ == "2.4.4"
    def bits(values): return np.asarray(values, dtype='<f8').view('<u8').tolist()
    for line in Path('tests/fixtures/mtplx-sampler-serial.jsonl').read_text().splitlines():
        case = json.loads(line)
        result = []
        for ids, probability_bits in zip(case['ids'], case['probability_bits']):
            distribution = _serial_row_distribution(np.asarray(ids, dtype=np.int64),
                np.asarray(probability_bits, dtype='<u8').view('<f8'), case['vocabulary'])
            if distribution is None:
                result.append(None)
                continue
            rng = np.random.default_rng(12345)
            samples = [sample_from_distribution(distribution, rng) for _ in range(32)]
            next_sample = bits(rng.random())
            draft = SparseDistribution(distribution.token_ids, distribution.probs[::-1], distribution.vocab_size)
            rng = np.random.default_rng(43210)
            decisions = []
            for _ in range(32):
                token = sample_from_distribution(draft, rng)
                decision = verify_one_token(distribution, draft, token, rng)
                decisions.append([token, decision.accepted, decision.token_id, bits(decision.accept_probability)])
            residual = residual_distribution(distribution, draft)
            result.append(dict(ids=distribution.token_ids.tolist(), probabilities=bits(distribution.probs),
                samples=samples, next_sample=next_sample, decisions=decisions, next_verify=bits(rng.random()),
                residual_ids=residual.token_ids.tolist(), residual_probabilities=bits(residual.probs)))
        print(json.dumps(result,separators=(',',':')),flush=True)


def sampler_serial_values(vocab, batch, mode):
    i = np.arange(batch*vocab, dtype=np.int64)
    unique = (i%vocab).astype(np.float32)/np.float32(vocab)
    sharp = -np.abs(i%vocab-vocab//2).astype(np.float32)*np.float32(0.2)
    if mode == "ties": return (i%7).astype(np.float32)
    if mode == "equal": return np.ones(batch*vocab, dtype=np.float32)
    if mode == "sharp": return sharp
    if mode == "nonfinite":
        unique[:3] = [np.nan, np.inf, -np.inf]
    if mode == "mixed":
        unique[:vocab] = 1
        if batch > 2: unique[2*vocab:3*vocab] = (i[2*vocab:3*vocab]%7).astype(np.float32)
        if batch > 3: unique[3*vocab:] = sharp[3*vocab:]
    return unique


def sampler_serial_fixtures():
    """Actual serial runtime helper with eval/selector calls recorded."""
    from mtplx import fast_sampling as fs
    from mtplx.sampling import SamplerConfig
    assert mx.__version__ == "0.32.2"
    assert np.__version__ == "2.4.4"
    original_eval = mx.eval
    original_selector = fs._deterministic_mlx_top_k_support
    counts = [0, 0]
    def evaluate(*args):
        counts[0] += 1
        return original_eval(*args)
    def selector(*args, **kwargs):
        counts[1] += 1
        return original_selector(*args, **kwargs)
    mx.eval = evaluate
    fs._deterministic_mlx_top_k_support = selector
    try:
        for vocab in (17, 4097, 248320):
            for batch in (1, 4):
                for mode in ("unique", "ties", "equal", "sharp", "nonfinite", "mixed"):
                    for top_p in (0.1, 0.95, 1.0):
                        temperature = 0.6
                        counts[:] = [0, 0]
                        logits = mx.array(sampler_serial_values(vocab, batch, mode)).reshape(batch, vocab)
                        ids, probabilities, size = fs._device_serial_support_arrays(logits, SamplerConfig(temperature=temperature, top_p=top_p))
                        probabilities = probabilities.astype('<f8')
                        probabilities[np.isnan(probabilities)] = np.nan
                        print(json.dumps(dict(vocab=vocab, batch=batch, mode=mode, top_p=top_p, temperature=temperature,
                            ids=ids.tolist(), probability_bits=probabilities.view('<u8').tolist(),
                            vocabulary=size, eval_calls=counts[0], selector_calls=counts[1]),separators=(',',':')),flush=True)
    finally:
        mx.eval = original_eval
        fs._deterministic_mlx_top_k_support = original_selector


def sampler_selection_fixtures():
    """The unchanged MTPLX deterministic selector, not a NumPy substitute."""
    from mtplx.fast_sampling import _deterministic_mlx_top_k_support
    assert mx.__version__ == "0.32.2"
    for vocab in (17, 4097, 248320):
        for batch in (1, 4):
            for mode in ("pattern", "ties", "equal", "nonfinite"):
                for top_k in (1, min(20, vocab)):
                    i = np.arange(batch*vocab, dtype=np.int64)
                    if mode == "pattern": values = ((i*17+96*13)%257-128).astype(np.float32)/128
                    elif mode == "ties": values = (i%7).astype(np.float32)
                    elif mode == "equal": values = np.ones(batch*vocab, dtype=np.float32)
                    else:
                        values = ((i*17+96*13)%257-128).astype(np.float32)/128
                        values[0] = np.nan
                        values[1] = np.inf
                        values[2] = -np.inf
                    rows = mx.array(values).reshape(batch, vocab)
                    ids, vals = _deterministic_mlx_top_k_support(rows, top_k)
                    emit(f"sampler_selection_v{vocab}_b{batch}_{mode}_k{top_k}",
                         [ids, vals], vocab=vocab, batch=batch, mode=mode, top_k=top_k)


def logsumexp_graph_fixtures():
    """Original sampler normalizer, including full Qwen vocabulary and layouts."""
    assert mx.__version__ == "0.32.2"
    for dtype in (mx.bfloat16, mx.float16, mx.float32):
        for cols in (1, 17, 4096, 4097, 248320):
            for layout in ("plain", "swapped", "strided", "offset", "broadcast"):
                if layout == "offset":
                    x = pattern(6*cols+1, 96).astype(dtype)[1:].reshape(2, 3, cols)
                elif layout == "strided":
                    x = pattern(12*cols, 96).astype(dtype).reshape(2, 3, cols*2)[..., 1::2]
                elif layout == "broadcast":
                    x = mx.broadcast_to(pattern(cols, 96).astype(dtype), (2, 3, cols))
                else:
                    x = pattern(6*cols, 96).astype(dtype).reshape(2, 3, cols)
                    if layout == "swapped":
                        x = x.transpose(1, 0, 2)
                for keepdims in (False, True):
                    name = str(dtype).split('.')[-1]
                    y = mx.logsumexp(x, axis=-1, keepdims=keepdims)
                    emit(f"logsumexp_{name}_{cols}_{layout}_{int(keepdims)}", [y],
                         dtype=name, cols=cols, layout=layout, keepdims=keepdims, shape=y.shape)


def softmax_graph_fixtures():
    """Original last-axis Softmax's precise/type/layout branches."""
    for dtype in (mx.bfloat16,mx.float16,mx.float32):
        for precise in (False,True):
            for cols in (17,4097):
                for layout in ("plain","swapped","strided","offset"):
                    if layout=="offset": x=pattern(6*cols+1,96).astype(dtype)[1:].reshape(2,3,cols)
                    elif layout=="strided": x=pattern(12*cols,96).astype(dtype).reshape(2,3,cols*2)[...,1::2]
                    else:
                        x=pattern(6*cols,96).astype(dtype).reshape(2,3,cols)
                        if layout=="swapped": x=x.transpose(1,0,2)
                    name=str(dtype).split('.')[-1]
                    emit(f"softmax_graph_{name}_{int(precise)}_{cols}_{layout}",[mx.softmax(x,axis=-1,precise=precise)],dtype=name,precise=precise,cols=cols,layout=layout)


def sdpa_fallback_fixtures():
    """Actual fast SDPA; these Qwen shapes select its original unfused graph."""
    assert mx.__version__ == "0.32.2"
    for batch,heads,rows,total in ((1,24,3,17),(1,24,4,1025),(1,24,8,4097),
                                  (1,24,16,33),(1,24,2048,2048),(2,24,3,17),(1,2,16,33)):
        q=pattern(batch*rows*heads*256,93).reshape(batch,rows,heads,256).transpose(0,2,1,3)
        k=pattern(batch*2*(total+7)*256,94).reshape(batch,2,total+7,256)[:,:,2:total+2,:]
        v=pattern(batch*2*(total+7)*256,95).reshape(batch,2,total+7,256)[:,:,2:total+2,:]
        for mode in ("none","bool","add"):
            visible=(np.arange(total)[None,:] <= (total-rows+np.arange(rows))[:,None]) & (np.arange(total)[None,:]%3!=1)
            visible[1,:]=False
            mask=None if mode=="none" else mx.array(visible).reshape(1,1,rows,total) if mode=="bool" else mx.array(np.where(visible,0.,-10.),mx.bfloat16).reshape(1,1,rows,total)
            out=mx.fast.scaled_dot_product_attention(q,k,v,scale=0.0625,mask=mask)
            emit(f"sdpa_fallback_B{batch}_h{heads}_r{rows}_t{total}_{mode}",[out],batch=batch,heads=heads,rows=rows,total=total,mode=mode)


def dense_batch_fixtures():
    """Original BF16 matmul, including Qwen's grouped attention layouts."""
    assert mx.__version__ == "0.32.2"
    for kind in ("qk","pv"):
        for batch,rows,total in ((1,1,17),(1,3,17),(1,8,17),(1,16,17),
                                 (1,64,1025),(1,2048,2048),(2,3,17),(2,16,33)):
            if kind=="qk":
                a=pattern(batch*rows*24*256,91).reshape(batch,rows,24,256).transpose(0,2,1,3).reshape(batch,2,12,rows,256)
                b=pattern(batch*2*total*256,92).reshape(batch,2,1,total,256).swapaxes(-1,-2)
            else:
                a=pattern(batch*24*rows*total,91).reshape(batch,2,12,rows,total)
                b=pattern(batch*2*total*256,92).reshape(batch,2,1,total,256)
            emit(f"dense_batch_{kind}_B{batch}_r{rows}_t{total}",[mx.matmul(a,b)],kind=kind,batch=batch,rows=rows,total=total)
    for kind in ("cross","transpose","copy","collapse","flatten","offset"):
        for rows in (1,3,16):
            k,n=64,37
            if kind=="cross":
                a=pattern(2*rows*k,91).reshape(2,1,rows,k)
                b=pattern(3*n*k,92).reshape(1,3,n,k).swapaxes(-1,-2)
            elif kind=="transpose":
                a=pattern(3*rows*k,91).reshape(3,k,rows).swapaxes(-1,-2)
                b=pattern(3*k*n,92).reshape(3,k,n)
            elif kind=="copy":
                a=pattern(3*rows*k*2,91).reshape(3,rows,k*2)[...,1::2]
                b=pattern(3*k*n*2,92).reshape(3,k,n*2)[...,1::2]
            elif kind in ("collapse","flatten"):
                a=pattern(3*rows*k,91).reshape(3,rows,k)
                b=pattern(n*k,92).reshape(n,k).T
                if kind=="collapse": b=b[None,...]
            else:
                a=pattern(3*rows*k+1,91)[1:].reshape(3,rows,k)
                b=pattern(3*n*k+1,92)[1:].reshape(3,n,k).swapaxes(-1,-2)
            emit(f"dense_batch_{kind}_r{rows}",[mx.matmul(a,b)],kind=kind,rows=rows)


def ple_projection_fixtures():
    """Actual PLELayer after supplied embedding: lookup is NOT tested here."""
    import mlx.nn as nn
    from mlx_lm.models.cache import ArraysCache
    from mtplx.models.qwen4_exp import PLELayer, TextArgs
    assert mx.__version__ == "0.32.2"
    class SuppliedEmbedding(nn.Module):
        def __call__(self, ids, cache, state_idx):
            return self.value
    for layout in ("dense", "quantized"):
        layer=PLELayer(TextArgs(ngram_sidecar=True),0)
        layer.ple_embedding=SuppliedEmbedding()
        for name,n,salt in (("key_proj",10240,71),("value_proj",2560,73)):
            if layout=="dense": getattr(layer,name).weight=pattern(n*2560,salt).reshape(n,2560)
            else: setattr(layer,name,quantized_linear_fixture(n,2560,4,64,salt))
        for name,salt in (("norm_key",75),("norm_query",76),("norm_conv",77)):
            getattr(layer,name).weight=pattern(10240,salt)
        layer.conv_weight=pattern(10240*4,78).reshape(10240,4,1)
        for batch,rows in ((1,1),(1,4),(1,7),(1,8),(1,32),(1,2048),(2,7)):
            for state in ("none","empty","seeded"):
                cache=None if state=="none" else ArraysCache(4)
                if state=="seeded": cache[2]=pattern(batch*12*10240,79).reshape(batch,12,10240)
                for step,tokens in enumerate((rows,1)):
                    hidden=pattern(batch*tokens*10240,80+step).reshape(batch,tokens,10240)
                    emb=pattern(batch*tokens*2560,82+step).reshape(batch,tokens,2560)
                    layer.ple_embedding.value=emb if step==0 else emb.astype(mx.float32)
                    output=layer(hidden,mx.zeros((batch,tokens),mx.int32),cache)
                    emit(f"ple_projection_{layout}_b{batch}_r{rows}_{state}_step{step}",
                         [output] if cache is None else [output,cache[2]],
                         layout=layout,batch=batch,rows=rows,state=state,step=step)


def shared_mlp_fixtures():
    """Only the real fused shared-expert class; no full model or downloads."""
    from mtplx.models.qwen4_exp import _FusedGateUpMLP
    assert mx.__version__ == "0.32.2"
    for rows,k,intermediate in ((1,128,64),(7,128,64),(1,2560,640),
                                (4,2560,640),(7,2560,640),(2048,2560,640)):
        for bits in (4,8):
            for group in (32,64):
                gu = quantized_linear_fixture(2*intermediate,k,bits,group,46)
                down = quantized_linear_fixture(k,intermediate,bits,group,48)
                layer = _FusedGateUpMLP(down,gu.weight,gu.scales,gu.biases,group,bits,"affine")
                layer.eval()
                x = pattern(rows*k,44).reshape(1,rows,k)
                emit(f"shared_mlp_r{rows}_k{k}_i{intermediate}_b{bits}_g{group}",
                     [layer(x)],rows=rows,k=k,intermediate=intermediate,bits=bits,group=group)


def gdn_fixture_layer():
    from mtplx.models.qwen4_exp import GatedDeltaNet, TextArgs, _FusedGDNInProj
    assert mx.__version__ == "0.32.2"
    for flag in ("MTPLX_FUSED_GDN_STEP", "MTPLX_FUSED_GDN_CONVNORM",
                 "MTPLX_FUSED_CONVNORM_VERIFY", "MTPLX_FUSED_GDN_OUT"):
        os.environ[flag] = "1"
    layer = GatedDeltaNet(TextArgs())
    projection = quantized_linear_fixture(16480, 2560, 4, 64, 26)
    layer.in_proj_fused = _FusedGDNInProj(projection.weight, projection.scales,
        projection.biases, 64, 4, "affine", [10240, 16384, 16432])
    layer.out_proj = quantized_linear_fixture(2560, 6144, 4, 64, 28)
    layer.conv1d.weight = pattern(10240*4, 8).reshape(10240,4,1)
    layer.A_log, layer.dt_bias, layer.norm.weight = pattern(48,12), pattern(48,13), pattern(128,14)
    layer.eval()
    return layer


def gdn_graph_lifecycle_fixtures():
    """Actual GatedDeltaNet: absent caches, batches and carried metadata/capture."""
    from mlx_lm.models.cache import ArraysCache
    from mtplx.models.qwen4_exp import _VERIFY_CAPTURE
    layer = gdn_fixture_layer()
    for batch in (1,2):
        for rows in (1,7):
            for mode in ("none", "plain", "metadata"):
                cache = None if mode == "none" else ArraysCache(2)
                if mode == "metadata":
                    cache.lengths = mx.array([rows+1,rows-1][:batch], dtype=mx.int32)
                    cache.left_padding = mx.array([1,0][:batch], dtype=mx.int32)
                mask = mx.array((np.arange(batch*rows)%3)!=0).reshape(batch,rows) if mode=="metadata" else None
                for step in range(2):
                    old_capture = getattr(cache, "_mtplx_verify_rows", None)
                    token = _VERIFY_CAPTURE.set(step == 0)
                    try:
                        output = layer(pattern(batch*rows*2560,30+step).reshape(batch,rows,2560),mask=mask,cache=cache)
                    finally:
                        _VERIFY_CAPTURE.reset(token)
                    outputs = [output] if cache is None else [output,cache[0],cache[1]]
                    if mode == "metadata":
                        outputs += [cache.lengths,cache.left_padding]
                    emit(f"gdn_lifecycle_b{batch}_r{rows}_{mode}_step{step}", outputs,
                         batch=batch,rows=rows,mode=mode,step=step,
                         capture_preserved=(getattr(cache,"_mtplx_verify_rows",None) is old_capture) if step else None)


def gdn_commit_fixtures():
    """Actual family commit method, with a pure-GDN layer and no model load."""
    from types import SimpleNamespace
    from mlx_lm.models.cache import ArraysCache
    from mtplx.models.qwen4_exp import Qwen4ExpTextModel, verify_capture_scope
    layer = gdn_fixture_layer()
    model = SimpleNamespace(layers=[SimpleNamespace(linear_attn=layer)])
    for rows in (1,2,4,6):
        for keep in range(1,rows+1):
            for conv_present in (False,True):
                cache = ArraysCache(2)
                cache[0] = pattern(3*10240,7).reshape(1,3,10240) if conv_present else None
                cache[1] = pattern(48*128*128,15).astype(mx.float32).reshape(1,48,128,128)
                snapshot = list(cache.state)
                with verify_capture_scope():
                    layer(pattern(rows*2560,30).reshape(1,rows,2560),cache=cache)
                mx.eval(cache.state, cache._mtplx_verify_rows)
                # Metadata is not rewound/advanced by the commit branch.
                cache.lengths = mx.array([17],mx.int32)
                cache.left_padding = mx.array([-3],mx.int32)
                assert Qwen4ExpTextModel.commit_verified_window(model,[cache],[snapshot],
                    keep_tokens=keep,verified_tokens=rows)
                assert cache._mtplx_verify_rows is None
                outputs = [cache[0],cache[1],cache.lengths,cache.left_padding]
                # Make the next decode consume the replay lazily. Mask/length
                # controls remain deliberately present as in the original cache.
                hidden = layer(pattern(2560,31).reshape(1,1,2560),cache=cache)
                outputs += [hidden,cache[0],cache[1]]
                emit(f"gdn_commit_r{rows}_keep{keep}_conv{int(conv_present)}", outputs,
                     rows=rows,keep=keep,conv_present=conv_present)


def qsa_f32_fixtures():
    """Run separately for MLX_ENABLE_TF32=0 and 1 (cached by the runtime).

    These are the actual MTPLX fallback expression, not reconstructed matmul
    microbenchmarks. Layouts cover initial/grown pooled mirrors and view copies.
    """
    from mtplx.kernels.qsa_indexer_prefill import qsa_indexer_prefill_scores
    assert mx.__version__ == "0.32.2"
    tf32 = os.environ.get("MLX_ENABLE_TF32", "1") != "0"
    for rows, blocks in ((1,1),(2,1),(4,7),(1,16),(4,32),(10,40),
                         (16,63),(16,64),(32,128),(33,129),(65,513),(2048,513)):
        layouts = [([512,128,1],[1,128]), ([512,1,4],[1,128]),
                   ([1024,256,2],[blocks+3,1]), ([519,128,1],[2,256]),
                   ([128,rows*128,1],[blocks,1]), ([512,1,4],[blocks+3,1]),
                   ([512,128,1],[blocks*2,2]), ([0,128,1],[1,128])]
        for layout, (qs, ps) in enumerate(layouts):
            for bf16 in (True, False):
                qcount = 1+sum((d-1)*s for d,s in zip((rows,4,128),qs))
                pcount = 1+(128-1)*ps[0]+(blocks-1)*ps[1]
                def backing(count, salt, bf16=False):
                    value = pattern(count,salt)
                    if not bf16:
                        value = value.astype(mx.float32)+mx.array(
                            (np.arange(count)%7).astype(np.float32)/65536)
                    return value
                q = mx.as_strided(backing(qcount,94,bf16), (1,rows,4,128), [0,*qs])
                p = mx.as_strided(backing(pcount,95), (1,1,128,blocks), [0,0,*ps])
                emit(f"qsa_f32_score_r{rows}_n{blocks}_l{layout}_bf{int(bf16)}_tf{int(tf32)}",
                     [qsa_indexer_prefill_scores(q,p,head_dim=128)], rows=rows,blocks=blocks,
                     q_strides=qs,pooled_strides=ps,q_count=qcount,pooled_count=pcount,
                     bf16=bf16,tf32=tf32)
    from mtplx.kernels.qsa_indexer_prefill import qsa_indexer_prefill_metal
    for rows,blocks,budget_rows in ((7,513,3),(65,600,48)):
        for ds in (1,2):
            for f32 in ((False,False),(False,True),(True,False),(True,True)):
                q = pattern(rows*512*ds,96)
                p = pattern(blocks*128*ds,97)
                q = q.astype(mx.float32) if f32[0] else q
                p = p.astype(mx.float32) if f32[1] else p
                # Cast before constructing strided F32 views, as in the Rust
                # fixture backing; otherwise AsType silently materializes them.
                q = q.reshape(1,rows,4,128*ds)[...,::ds]
                p = p.reshape(1,blocks,128*ds)[...,::ds]
                total = blocks*4+1
                budget = blocks*4*5*budget_rows
                for mode in ("blocks","dense_mask","row_tokens"):
                    out = qsa_indexer_prefill_metal(q,p,pos_start=total-rows,total_tokens=total,
                        block_topk=512,compress_ratio=4,logical_blocks=blocks,
                        output_total_tokens=total+3,mode=mode,score_workspace_bytes=budget)
                    emit(f"qsa_f32_chain_r{rows}_ds{ds}_f{int(f32[0])}{int(f32[1])}_tf{int(tf32)}_{mode}",
                         [out] if mode=="dense_mask" else out,rows=rows,blocks=blocks,ds=ds,
                         input_f32=f32,tf32=tf32,total=total,budget=budget,mode=mode)


def qsa_eager_fixtures():
    """Tap the real eager oracle without replacing its arithmetic or sorting."""
    from mtplx.models.qwen4_exp import QSAIndexer, QSACache, TextArgs
    from mtplx.attention_context import attention_phase
    os.environ["MTPLX_QSA_SCORE_TILE_ROWS"] = "0"
    tf32 = os.environ.get("MLX_ENABLE_TF32", "1") != "0"
    indexer = QSAIndexer(TextArgs())
    for rows,blocks in ((1,1),(4,7),(1,513),(4,513),(7,600),
                        (33,2049),(65,513),(2048,513),(2048,1025),(1,65536)):
        total = blocks*4+3
        for variant in range(3):
            q = pattern(rows*512,98).reshape(1,rows,4,128)
            if variant == 1: q = mx.zeros_like(q)
            if variant == 2: q = q.astype(mx.float32)
            pooled = pattern(blocks*128,99).reshape(1,blocks,128)
            cache = QSACache()
            cache.pooled,cache.pooled_len = pooled,blocks
            original_partition, original_where = mx.argpartition,mx.where
            tapped = {}
            def partition(a,*args,**kwargs):
                result = original_partition(a,*args,**kwargs)
                assert "sorted" not in tapped, "unexpected tiled eager oracle"
                tapped.update(masked=a,sorted=result)
                return result
            def where(condition,a,b,*args,**kwargs):
                if "valid" not in tapped:
                    tapped["valid"] = condition
                return original_where(condition,a,b,*args,**kwargs)
            try:
                mx.argpartition,mx.where = partition,where
                indexer._select_eager(q,total-rows,cache,pooled,total)
            finally:
                mx.argpartition,mx.where = original_partition,original_where
            emit(f"qsa_eager_topk_r{rows}_n{blocks}_v{variant}_tf{int(tf32)}",
                 [tapped["sorted"][:,blocks-min(blocks,512):],tapped["valid"],tapped["masked"]],
                 rows=rows,blocks=blocks,variant=variant,total=total,tf32=tf32)
            if rows > 1 and total-rows >= 2049:
                knobs = {"MTPLX_QSA_PREFILL":"1","MTPLX_QSA_PREFILL_MIN_ROWS":"2",
                         "MTPLX_QSA_PREFILL_MIN_CONTEXT":"2049"}
                saved = {k:os.environ.get(k) for k in knobs}
                try:
                    os.environ.update(knobs)
                    with attention_phase("prefill"):
                        result = indexer._select_eager(q,total-rows,cache,pooled,total)
                finally:
                    for key,value in saved.items():
                        if value is None: os.environ.pop(key,None)
                        else: os.environ[key] = value
                assert result[0] == "flash_prefill"
                emit(f"qsa_eager_blocks_r{rows}_n{blocks}_v{variant}_tf{int(tf32)}",result[1:],
                     rows=rows,blocks=blocks,variant=variant,total=total,tf32=tf32)
            if blocks > 512:
                knobs = {"MTPLX_QSA_PREFILL":"0","MTPLX_QSA_FLASH":"0",
                         "MTPLX_QSA_GATHER":"0","MTPLX_QSA_GATHER_DECODE":"0",
                         "MTPLX_QSA_GATHER_MIN_CONTEXT":"0","MTPLX_QSA_GATHER_MAX_ROWS":str(max(rows,2))}
                saved = {k:os.environ.get(k) for k in knobs}
                try:
                    os.environ.update(knobs)
                    for tail in (0,1,3):
                        dense_total = blocks*4+tail
                        result = indexer._select_eager(q,dense_total-rows,cache,pooled,dense_total)
                        assert isinstance(result,mx.array) and result.dtype == mx.bool_
                        emit(f"qsa_eager_dense_r{rows}_n{blocks}_v{variant}_tail{tail}_tf{int(tf32)}",[result],
                             rows=rows,blocks=blocks,variant=variant,total=dense_total,tf32=tf32)
                        if rows > 1:
                            os.environ["MTPLX_QSA_GATHER"] = "1"
                            result = indexer._select_eager(q,dense_total-rows,cache,pooled,dense_total)
                            os.environ["MTPLX_QSA_GATHER"] = "0"
                            assert result[0] == "gather_rows"
                            emit(f"qsa_eager_gather_r{rows}_n{blocks}_v{variant}_tail{tail}_tf{int(tf32)}",result[1:],
                                 rows=rows,blocks=blocks,variant=variant,total=dense_total,tf32=tf32)
                        else:
                            for gather in (False,True):
                                os.environ["MTPLX_QSA_FLASH"] = "0" if gather else "1"
                                os.environ["MTPLX_QSA_GATHER_DECODE"] = "1"
                                result = indexer._select_eager(q,dense_total-1,cache,pooled,dense_total)
                                if gather:
                                    assert isinstance(result,mx.array)
                                    values,tail_start = [result],dense_total//4*4
                                else:
                                    assert result[0] == "flash"
                                    values,tail_start = [result[1]],result[2]
                                emit(f"qsa_eager_decode_r1_n{blocks}_v{variant}_tail{tail}_g{int(gather)}_tf{int(tf32)}",values,
                                     rows=1,blocks=blocks,variant=variant,total=dense_total,tf32=tf32,gather=gather,tail_start=tail_start)
                            os.environ["MTPLX_QSA_FLASH"] = "0"
                            os.environ["MTPLX_QSA_GATHER_DECODE"] = "0"
                finally:
                    for key,value in saved.items():
                        if value is None: os.environ.pop(key,None)
                        else: os.environ[key] = value


def qsa_eager_tiled_fixtures():
    """Actual eager entry point, including per-tile eval and output routing."""
    from mtplx.models.qwen4_exp import QSAIndexer, TextArgs, QSACache
    from mtplx.attention_context import attention_phase
    tf32 = os.environ.get("MLX_ENABLE_TF32", "1") != "0"
    indexer = QSAIndexer(TextArgs())
    knobs = {"MTPLX_QSA_SCORE_TILE_ROWS":"0", "MTPLX_QSA_PREFILL":"0",
             "MTPLX_QSA_PREFILL_MIN_ROWS":"2", "MTPLX_QSA_PREFILL_MIN_CONTEXT":"2049",
             "MTPLX_QSA_GATHER":"1", "MTPLX_QSA_GATHER_MIN_CONTEXT":"0",
             "MTPLX_QSA_GATHER_MAX_ROWS":"2048", "MTPLX_QSA_FLASH":"1",
             "MTPLX_QSA_GATHER_DECODE":"1"}
    saved = {k:os.environ.get(k) for k in knobs}
    try:
        os.environ.update(knobs)
        for rows,blocks,tiles in ((7,513,(1,3,7,8)), (33,600,(3,32)),
                                 (65,2049,(1,32)), (2048,1025,(32,2047)),
                                 (7,65536,(3,))):
            for variant in range(3):
                ds = 2 if variant == 2 else 1
                q = pattern(rows*512*ds,98)
                if variant == 1: q = mx.zeros_like(q)
                if variant == 2: q = q.astype(mx.float32)
                q = q.reshape(1,rows,4,128*ds)[...,::ds]
                pooled = pattern(blocks*128,99).reshape(1,blocks,128)
                cache = QSACache()
                cache.pooled,cache.pooled_len = pooled,blocks
                for tile in tiles:
                    total = blocks*4+variant%2
                    os.environ["MTPLX_QSA_SCORE_TILE_ROWS"] = str(tile)
                    for prefill in (False,True):
                        os.environ["MTPLX_QSA_PREFILL"] = str(int(prefill))
                        eval_rows = []
                        original_eval = mx.eval
                        def evaluate(*arrays):
                            assert len(arrays) == 1 and arrays[0].dtype == mx.uint32
                            eval_rows.append(arrays[0].shape[0])
                            return original_eval(*arrays)
                        try:
                            mx.eval = evaluate
                            with attention_phase("prefill"):
                                result = indexer._select_eager(q,total-rows,cache,pooled,total)
                        finally:
                            mx.eval = original_eval
                        expected = [min(tile,rows-start) for start in range(0,rows,tile)] if tile < rows else []
                        assert eval_rows == expected, (eval_rows,expected)
                        if prefill and total-rows >= 2049:
                            assert result[0] == "flash_prefill"
                            kind, values = "blocks", result[1:]
                        elif tile >= rows:
                            assert result[0] == "gather_rows"
                            kind, values = "gather", result[1:]
                        else:
                            # Tiled scoring disables rows-gather even if all
                            # its flags and thresholds would otherwise engage.
                            assert isinstance(result,mx.array) and result.dtype == mx.bool_
                            kind, values = "dense", [result]
                        emit(f"qsa_eager_{kind}_tiled_r{rows}_n{blocks}_v{variant}_t{tile}_p{int(prefill)}_tf{int(tf32)}",
                             values,rows=rows,blocks=blocks,variant=variant,total=total,tf32=tf32,
                             tile=tile,ds=ds,prefill=prefill,eval_rows=eval_rows)
    finally:
        for key,value in saved.items():
            if value is None: os.environ.pop(key,None)
            else: os.environ[key] = value


def qsa_compiled_state_fixtures():
    """Original compiled backing policy and device-indexed copy primitives.

    This checks prerequisites, not execution of the compiled indexer graph.
    """
    from mtplx.qsa_mtp_precompute import precompute_qsa_replay_capacity
    from mtplx.models.qwen4_exp import QSAIndexer, TextArgs, QSACache
    for start in (0,3,2048,4095):
        for window in (0,1,3,4,255,256,257,1025,2048):
            for raw,pool in ((0,0),(768,256),(8192,4096)):
                plan = precompute_qsa_replay_capacity(start_offset=start,window_tokens=window,
                    compress_ratio=4,current_raw_capacity=raw,current_pooled_capacity=pool)
                emit(f"qsa_compile_capacity_p{start}_s{window}_c{raw}_{pool}",[],
                    start=start,window=window,existing=[raw,pool],capacity=list(plan.graph_key))
    indexer,cache = QSAIndexer(TextArgs()),QSACache()
    for step,(rows,reserve) in enumerate(((1,(0,0)),(1025,(0,0)),(2048,(8192,4096)),(8193,(0,0)))):
        if reserve != (0,0):
            cache.reserve_indexer_capacity(raw_capacity=reserve[0],pooled_capacity=reserve[1])
        raw,pool = indexer._ensure_compiled_backings(cache,dtype=mx.bfloat16,pos_start=0,rows=rows)
        emit(f"qsa_compile_materialize_{step}",[raw,pool],step=step,rows=rows,reserve=reserve,
            capacity=[raw.shape[1],pool.shape[1]],offset=cache.offset,pooled_len=cache.pooled_len)
        # Make the next capacity transition prove preservation, not just zero fill.
        cache.write_raw(pattern(3*128,160+step).reshape(1,3,128))
        mx.eval(cache.raw_keys)
    for rows in (1,4,33,2048):
        for stride in (128,704):
            cap = rows+7
            source = pattern(cap*stride,151).reshape(1,cap,stride)[...,:128]
            update = pattern(rows*stride,152).reshape(1,rows,stride)[...,:128]
            start = mx.array([3],dtype=mx.int32)
            mx.eval(source,update,start)
            sliced = mx.slice(source,start,axes=(1,),slice_size=(1,rows,128))
            updated = mx.slice_update(source,update,start,axes=(1,))
            emit(f"qsa_compile_dynamic_r{rows}_stride{stride}",[sliced,updated,source],
                rows=rows,stride=stride,capacity=cap,start=3)


def qsa_compiled_window_fixtures():
    """Real compiled cache-maintenance calls with retained original leaves."""
    from mtplx.models.qwen4_exp import QSAIndexer, TextArgs
    from mtplx.kernels.qsa_indexer_compile import QSACompiledIndexerCore
    freq = QSAIndexer(TextArgs())._inv_freq
    norm = pattern(128,72)
    cases = ((1,256,256,0,0),(1,256,256,253,63),(1,256,256,255,63),
        (4,4096,1024,2051,512),(5,4096,1024,2050,511),(5,256,256,0,0),
        (2048,2048,512,0,0),(2048,4096,1024,2048,512),
        (1025,2048,512,100,24),(1,4096,1024,4095,1024),
        (33,4096,1024,2050,512),(1020,1024,256,0,0),(512,1024,256,0,0))
    for case,(rows,raw_cap,pool_cap,pos,pool_len) in enumerate(cases):
        core = QSACompiledIndexerCore(n_heads=4,kv_heads=1,head_dim=128,block_topk=512,
            compress_ratio=4,q_norm_weight=norm,k_norm_weight=norm,inv_freq=freq,rms_norm_eps=1e-6)
        for stride in (640,704):
            qk = pattern(rows*stride,150).reshape(1,rows,stride)[...,:640]
            raw = pattern(raw_cap*128,151).reshape(1,raw_cap,128)
            pool = pattern(pool_cap*128,152).reshape(1,pool_cap,128)
            total = pos+rows
            logical = total//4
            scalars = [mx.array([n],dtype=mx.int32) for n in (pos,total,logical,pool_len)]
            mx.eval(qk,raw,pool,norm,freq,*scalars)
            result = core.select_qk_rows(qk,raw,pool,pos_start=scalars[0],total_tokens=scalars[1],
                logical_blocks=scalars[2],pooled_len=scalars[3],mode="update_only")
            assert result.selection is None
            # Keep raw/pool live across eval: this is the reference copy branch,
            # not evidence for donation or the complete model's lazy schedule.
            emit(f"qsa_compile_window_c{case}_stride{stride}",[result.raw_keys,result.pooled,raw,pool],
                rows=rows,stride=stride,capacity=[raw_cap,pool_cap],pos=pos,logical=logical,
                pooled_len=pool_len,frequencies=freq.tolist())
            assert result.pooled_len.tolist()==[logical] and result.offset.tolist()==[total]


def qsa_compiled_graph_contract(result, selection, inputs):
    """Fingerprint the actual optimized runtime graph before evaluation.

    Evaluated constant leaves are omitted, not pending primitive dependencies.
    Multi-output leaf positions remain explicit. No execution-order inference
    is made from the DOT traversal order.
    """
    import io
    names = dict(inputs)
    roots = [f"selected{i}" for i in range(len(selection))] + ["raw_next","pool_next"]
    names.update({f"selected{i}":v for i,v in enumerate(selection)})
    names.update(raw_next=result.raw_keys,pool_next=result.pooled)
    dot = io.StringIO()
    mx.export_to_dot(dot,**names)
    return compiled_dot_contract(dot.getvalue(), roots, inputs)


def compiled_dot_contract(dot, roots, inputs):
    """Shared ordered-dependency fingerprint for an original runtime DOT."""
    import re
    nodes,producer = {},{}
    for line in dot.splitlines():
        if m:=re.search(r'(\d+) \[label ="([^"]+)"',line):
            nodes[m[1]] = dict(op=m[2],inputs=[],outputs=[])
        elif m:=re.fullmatch(r'"([^"]+)" -> (\d+)',line):
            nodes[m[2]]["inputs"].append(m[1])
        elif m:=re.fullmatch(r'(\d+) -> "([^"]+)"',line):
            node = nodes[m[1]]
            producer[m[2]] = (m[1],len(node["outputs"]))
            node["outputs"].append(m[2])
    hashes = {}
    def leaf(name):
        if name in producer:
            node,index = producer[name]
            return fingerprint(node)+":"+str(index)
        return name if name in inputs else None
    def fingerprint(key):
        if key not in hashes:
            node = nodes[key]
            parents = [value for name in node["inputs"] if (value:=leaf(name)) is not None]
            hashes[key] = hashlib.sha256(json.dumps([node["op"],parents],separators=(",",":")).encode()).hexdigest()
        return hashes[key]
    root_hashes = [leaf(name) for name in roots]
    return [sorted(hashes.values()),root_hashes]


def qsa_compiled_indexer_fixtures(*,hidden=False,bits=4,group=64,graph=False):
    """All fixed-output modes of the original compiled BF16 Q/K or hidden entry."""
    from mtplx.models.qwen4_exp import QSAIndexer, TextArgs
    from mtplx.kernels.qsa_indexer_compile import QSACompiledIndexerCore
    freq = QSAIndexer(TextArgs())._inv_freq
    norm = pattern(128,72)
    projection = None
    if hidden:
        import mlx.nn as nn
        projection = nn.QuantizedLinear(2560,640,bias=False,group_size=group,bits=bits)
        words = (np.arange(640*2560*bits//32,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
        projection.weight = mx.array(words).reshape(640,2560*bits//32)
        projection.scales = (pattern(640*2560//group,75)/64).reshape(640,2560//group)
        projection.biases = (pattern(640*2560//group,76)/64).reshape(640,2560//group)
    modes = ("blocks","row_tokens","dense_mask")
    scenarios = [(rows,pos,frontier,mode,budget) for rows,pos,frontier in
        ((1,2051,512),(4,2051,512),(5,2050,511),(33,2050,512))
        for mode in modes for budget in (32*1024*1024,3*1024*4)]
    scenarios += [(rows,pos,frontier,"prefill_blocks",budget) for rows,pos,frontier in
        ((4,2051,512),(33,2050,512),(65,2050,512),(2048,2048,512))
        for budget in (128*1024*1024,(128 if rows==2048 else 48 if rows==65 else 3)*1024*4)]
    scenarios += [(4,2051,512,"update_only",32*1024*1024)]
    for case,(rows,pos,pool_len,mode,budget) in enumerate(scenarios):
        core = QSACompiledIndexerCore(n_heads=4,kv_heads=1,head_dim=128,block_topk=512,
            compress_ratio=4,q_norm_weight=norm,k_norm_weight=norm,inv_freq=freq,rms_norm_eps=1e-6,
            selector_scratch_bytes=budget,prefill_score_workspace_bytes=budget,project_qk=projection)
        for stride in ((2560,) if hidden else (640,704)):
            source = pattern(rows*stride,150).reshape(1,rows,stride)
            if not hidden: source = source[...,:640]
            raw = pattern(4096*128,151).reshape(1,4096,128)
            pool = pattern(1024*128,152).reshape(1,1024,128)
            total,logical = pos+rows,(pos+rows)//4
            state = [mx.array([n],dtype=mx.int32) for n in (pos,total,logical,pool_len)]
            mx.eval(source,raw,pool,norm,freq,*state)
            if graph and projection is not None:
                mx.eval(projection.weight,projection.scales,projection.biases)
            entry = core.select_hidden if hidden else core.select_qk_rows
            out = entry(source,raw,pool,pos_start=state[0],total_tokens=state[1],
                logical_blocks=state[2],pooled_len=state[3],mode=mode)
            selection = [] if out.selection is None else list(out.selection) if isinstance(out.selection,tuple) else [out.selection]
            if graph:
                contract = qsa_compiled_graph_contract(out,selection,dict(source=source,raw=raw,pooled=pool,
                    norm=norm,freq=freq,pos=state[0],total=state[1],logical=state[2],frontier=state[3]))
                print(json.dumps(dict(case=case,hidden=hidden,rows=rows,stride=stride,mode=mode,budget=budget,
                    contract=contract),separators=(",",":")),flush=True)
                mx.eval(out.raw_keys,out.pooled,selection)
                continue
            mx.eval(out.raw_keys,out.pooled,selection)
            states = {"raw_next":out.raw_keys,"pool_next":out.pooled,"raw_old":raw,"pool_old":pool}
            hashes = {key:hashlib.sha256(np.asarray(a.astype(mx.float32)).astype('<f4').tobytes()).hexdigest() for key,a in states.items()}
            prefix = f"qsa_compile_hidden_b{bits}_g{group}" if hidden else "qsa_compile_indexer"
            extra = dict(bits=bits,group=group) if hidden else {}
            emit(f"{prefix}_c{case}_stride{stride}",selection,rows=rows,stride=stride,mode=mode,
                budget=budget,capacity=[4096,1024],pos=pos,logical=logical,pooled_len=pool_len,
                frequencies=freq.tolist(),state_sha256=hashes,**extra)


def qsa_compiled_route_fixtures():
    """Original host routing methods, with real typed arrays; no graph eval.

    This covers the installed B1/BF16 contract, not unsupported tensor dtypes.
    Cache extents deliberately include non-power-of-two eager backings.
    """
    from mtplx.models.qwen4_exp import QSAIndexer, TextArgs, QSACache
    from mtplx.attention_context import attention_phase
    assert mx.__version__ == "0.32.2"
    assert mx.metal.is_available() and mx.default_device() == mx.gpu
    defaults = dict(fused=True,compiled=True,prefill=True,phase_prefill=True,
        prefill_min_rows=32,prefill_compile_rows=2048,prefill_min_context=32768,
        prefill_workspace=128*1024*1024,tile_rows=0,flash=True,gather_decode=True,
        gather=True,gather_max_rows=8,gather_min_context=0,tf32=True)
    env_names = dict(fused="MTPLX_FUSED_QSA_INDEXER",compiled="MTPLX_COMPILED_QSA_INDEXER",
        prefill="MTPLX_QSA_PREFILL",prefill_min_rows="MTPLX_QSA_PREFILL_MIN_ROWS",
        prefill_compile_rows="MTPLX_QSA_PREFILL_COMPILE_ROWS",
        prefill_min_context="MTPLX_QSA_PREFILL_MIN_CONTEXT",flash="MTPLX_QSA_FLASH",
        gather_decode="MTPLX_QSA_GATHER_DECODE",gather="MTPLX_QSA_GATHER",
        gather_max_rows="MTPLX_QSA_GATHER_MAX_ROWS",gather_min_context="MTPLX_QSA_GATHER_MIN_CONTEXT")
    options = [defaults | change for change in ({},dict(fused=False),dict(compiled=False),
        dict(prefill=False),dict(phase_prefill=False),dict(flash=False),
        dict(flash=False,gather_decode=False),dict(gather=False),dict(gather_max_rows=3),
        dict(gather_min_context=32772),dict(prefill_compile_rows=32),
        dict(prefill_min_rows=2,prefill_min_context=2049))]
    # pos, rows, offset, raw capacity, pooled capacity, pooled frontier
    states = []
    for pos,rows in ((0,0),(0,1),(0,3),(0,32),(0,2048),(2048,1),(2048,4),
                     (2048,2048),(2049,4),(2051,1),(2051,4),(2051,8),(2051,9),
                     (32767,2048),(32768,1),(32768,31),(32768,32),(32768,33),
                     (32768,2047),(32768,2048),(32768,2049)):
        states.append([pos,rows,pos,((pos+255)//256)*256 if pos else None,
                       ((pos//4+255)//256)*256 if pos else None,pos//4])
    states.extend([
        [1,1,0,256,256,0],       # mismatched offset
        [1,1,1,None,None,0],      # missing prior raw
        [257,1,257,256,256,64],   # short raw
        [0,1,0,None,None,1],      # nonzero frontier without a pool
        [2051,4,2051,2304,512,513], # frontier exceeds capacity
        [2051,4,2051,2304,768,511], # missing two blocks, window fits one
        [2051,4,2051,2304,768,512], # valid nonbucketed backing
        [2051,4,2051,2304,768,700], # frontier beyond logical: clamp first
        [0,1,0,0,0,0],           # present zero-length leaves
    ])
    saved = {env:os.environ.get(env) for env in env_names.values()}
    indexer = QSAIndexer(TextArgs())
    indexer.q_layernorm.weight = mx.ones((128,),mx.bfloat16)
    indexer.k_layernorm.weight = mx.ones((128,),mx.bfloat16)
    indexer.index_qk_proj.weight = mx.zeros((640,2560),mx.bfloat16)
    cases = []
    try:
        for option_id,opt in enumerate(options):
            os.environ.update({env:str(int(opt[key])) for key,env in env_names.items()})
            with attention_phase("prefill" if opt["phase_prefill"] else "verify"):
                for state_id,(pos,rows,offset,raw,pool,frontier) in enumerate(states):
                    cache = QSACache()
                    cache.kv.offset = offset
                    cache.raw_keys = None if raw is None else mx.zeros((1,raw,128),mx.bfloat16)
                    cache.pooled = None if pool is None else mx.zeros((1,pool,128),mx.bfloat16)
                    cache.pooled_len = frontier
                    mode = indexer._compiled_mode(decode=rows==1,rows=rows,
                        total=pos+rows,last_nb=(pos+rows)//4)
                    results = []
                    for supplied in (False,True):
                        source = mx.zeros((1,rows,640 if supplied else 2560),mx.bfloat16)
                        results.append(indexer._compiled_route_supported(source,cache,
                            pos_start=pos,qk_rows_supplied=supplied,decode=rows==1,mode=mode))
                    cases.append([option_id,state_id,mode,*results])
    finally:
        for env,value in saved.items():
            if value is None: os.environ.pop(env,None)
            else: os.environ[env] = value
    print(json.dumps(dict(options=options,states=states,cases=cases),separators=(",",":")),flush=True)


def qsa_indexer_route_fixtures(*,compiled=False):
    """Whole original indexer calls, not manually selected outputs."""
    from mtplx.models.qwen4_exp import QSAIndexer, TextArgs, QSACache
    from mtplx.attention_context import attention_phase
    import mlx.nn as nn
    defaults = dict(fused=False,prefill=False,phase_prefill=True,prefill_min_rows=32,
        prefill_min_context=2049,prefill_workspace=128*1024*1024,tile_rows=0,
        flash=False,gather_decode=False,gather=False,gather_max_rows=8,
        gather_min_context=0,tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0")
    env_names = dict(fused="MTPLX_FUSED_QSA_INDEXER",prefill="MTPLX_QSA_PREFILL",
        prefill_min_rows="MTPLX_QSA_PREFILL_MIN_ROWS",prefill_min_context="MTPLX_QSA_PREFILL_MIN_CONTEXT",
        tile_rows="MTPLX_QSA_SCORE_TILE_ROWS",flash="MTPLX_QSA_FLASH",
        gather_decode="MTPLX_QSA_GATHER_DECODE",gather="MTPLX_QSA_GATHER",
        gather_max_rows="MTPLX_QSA_GATHER_MAX_ROWS",gather_min_context="MTPLX_QSA_GATHER_MIN_CONTEXT")
    knobs = [*env_names.values(),"MTPLX_COMPILED_QSA_INDEXER","MTPLX_QSA_PREFILL_SCORE_MB",
             "MTPLX_QSA_PREFILL_COMPILE_ROWS"]
    saved = {k:os.environ.get(k) for k in knobs}
    scenarios = [
        (0,(1,3,4,33,2048),{}), (0,(1,3,4,33,2048),dict(fused=True)),
        (2048,(1,1,1,1,4,33),dict(flash=True,gather=True,gather_min_context=2053)),
        (2048,(1,1,1,1,4,33),dict(fused=True,flash=True,gather=True,gather_min_context=2053)),
        (2051,(1,4,33),dict(fused=True,flash=True,gather_decode=True,gather=True)),
        (2051,(1,4,33),dict(fused=True,gather_decode=True)),
        (2051,(1,4,33),dict(fused=True,gather=True,gather_max_rows=3)),
        (2051,(4,33),dict(gather=True,tile_rows=3)),
        (2048,(32,33,2048),dict(fused=True,prefill=True)),
        (2048,(32,33,2048),dict(fused=True,prefill=True,phase_prefill=False)),
        (32768,(4,32,33,2048),dict(prefill=True,prefill_min_context=32768,gather=True,gather_min_context=16384)),
        (2051,(1,4,33),dict(gather=True,gather_min_context=99999)),
    ]
    try:
        os.environ["MTPLX_COMPILED_QSA_INDEXER"] = str(int(compiled))
        os.environ["MTPLX_QSA_PREFILL_SCORE_MB"] = "128"
        os.environ["MTPLX_QSA_PREFILL_COMPILE_ROWS"] = "2048"
        if compiled:
            scenarios.extend([
                (32768,(2048,33,1,4),dict(fused=True,prefill=True,flash=True,gather=True)),
                (0,(1,2048,1,4),dict(fused=True,phase_prefill=False)),
            ])
        for scenario,(prefix,steps,overrides) in enumerate(scenarios):
            options = defaults | overrides
            if compiled: options.update(compiled=True,prefill_compile_rows=2048)
            os.environ.update({env:str(int(options[key])) for key,env in env_names.items()})
            for supplied in (False,True):
                indexer = QSAIndexer(TextArgs())
                indexer.q_layernorm.weight = pattern(128,72)
                indexer.k_layernorm.weight = pattern(128,73)
                proj = nn.QuantizedLinear(2560,640,bias=False,group_size=64,bits=4)
                words = (np.arange(640*2560*4//32,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
                proj.weight = mx.array(words).reshape(640,2560*4//32)
                proj.scales = (pattern(640*2560//64,75)/64).reshape(640,2560//64)
                proj.biases = (pattern(640*2560//64,76)/64).reshape(640,2560//64)
                indexer.index_qk_proj = proj
                cache = QSACache()
                if prefix:
                    cache.write_raw(pattern(prefix*128,88).reshape(1,prefix,128))
                    indexer._extend_pooled(cache,prefix)
                    kv = pattern(prefix*512,89).reshape(1,prefix,2,256).transpose(0,2,1,3)
                    cache.kv.update_and_fetch(kv,kv)
                    mx.eval(*cache.state,cache.pooled_f32_t)
                lane = ["update_only"]
                eager,fused,prefill = indexer._select_eager,indexer._select_fused,indexer._prefill_selector_supported
                def eager_call(*args,**kwargs):
                    lane[0] = "eager"
                    return eager(*args,**kwargs)
                def fused_call(*args,**kwargs):
                    lane[0] = "legacy_fused"
                    return fused(*args,**kwargs)
                def prefill_call(*args,**kwargs):
                    result = prefill(*args,**kwargs)
                    if result: lane[0] = "metal_prefill"
                    return result
                indexer._select_eager,indexer._select_fused,indexer._prefill_selector_supported = eager_call,fused_call,prefill_call
                original_compiled = indexer._call_rows_compiled
                def compiled_call(*args,**kwargs):
                    lane[0] = "compiled"
                    return original_compiled(*args,**kwargs)
                indexer._call_rows_compiled = compiled_call
                for step,rows in enumerate(steps):
                    pos = cache.offset
                    hidden = pattern(rows*2560,77+step).reshape(1,rows,2560)
                    stride = 704 if supplied else 640
                    qk = pattern(rows*stride,121+step).reshape(1,rows,stride)[...,:640] if supplied else None
                    lane[0] = "update_only"
                    with attention_phase("prefill" if options["phase_prefill"] else "verify"):
                        result = indexer(hidden,pos,cache,qk)
                    tail = None
                    if result is None: kind,values = "none",[]
                    elif isinstance(result,tuple):
                        kind = result[0]
                        if kind == "flash": values,tail = [result[1]],result[2]
                        else: values = list(result[1:])
                    else: kind,values = ("dense" if result.dtype == mx.bool_ else "decode_tokens"),[result]
                    states = {"raw":cache.raw_keys,"pooled":cache.pooled,"mirror":cache.pooled_f32_t}
                    mx.eval(*values,*(a for a in states.values() if a is not None))
                    state_hashes = {key:None if a is None else hashlib.sha256(np.asarray(a.astype(mx.float32)).astype('<f4').tobytes()).hexdigest() for key,a in states.items()}
                    tag = "qsa_indexer_compiled" if compiled else "qsa_indexer"
                    graph_counts = {}
                    if compiled:
                        core = indexer._compiled_indexer_core
                        graph_counts["graph_counts"] = [0,0,0] if core is None else [
                            core.stats["compiled_calls"],core.stats["traces"],len(core._compiled)]
                    emit(f"{tag}_s{scenario}_qk{int(supplied)}_step{step}",values,
                        scenario=scenario,step=step,prefix=prefix,rows=rows,pos=pos,total=pos+rows,
                        supplied=supplied,stride=stride,options=options,kind=kind,lane=lane[0],tail=tail,
                        raw_cap=cache.raw_keys.shape[1],pooled_cap=0 if cache.pooled is None else cache.pooled.shape[1],
                        pooled_len=cache.pooled_len,mirror_cap=0 if cache.pooled_f32_t is None else cache.pooled_f32_t.shape[3],
                        state_sha256=state_hashes,frequencies=np.asarray(indexer._inv_freq).tolist(),**graph_counts)
                    assert cache.offset == pos
                    kv = pattern(rows*512,140+step).reshape(1,rows,2,256).transpose(0,2,1,3)
                    cache.kv.update_and_fetch(kv,kv)
                    mx.eval(*cache.kv.state)
    finally:
        for key,value in saved.items():
            if value is None: os.environ.pop(key,None)
            else: os.environ[key] = value


def main():
    assert mx.__version__ == "0.32.2", mx.__version__
    with (Path(mx.__file__).parent / "lib/mlx.metallib").open("rb") as library:
        assert hashlib.file_digest(library, "sha256").hexdigest() == (
            "dc59d1cceb1a5c7e578232e6e41e28e2c73c9463ac6dbc3886c3ee17ffc270ed")
    reference = Path(hyper_connection.__file__).resolve().parents[2]
    exporter = runpy.run_path(str(Path(__file__).with_name("mtplx-kernel-source.py")))
    assert exporter["generate"](reference, Path(gated_delta.__file__)) == exporter["OUTPUT"].read_text()
    assert exporter["qsa_select_source"](reference) == exporter["QSA_SELECT_OUTPUT"].read_text()
    print(f"MTPLX {exporter['REVISION']}, runtime {mx.__version__}, {reference}", file=sys.stderr)
    qsa_f32_fixtures()
    qsa_eager_fixtures()
    qsa_eager_tiled_fixtures()
    qsa_indexer_route_fixtures()
    for rows in (1, 4, 8):
        x = pattern(rows * 10240, 1).reshape(rows, 10240)
        gamma = pattern(10240, 2)
        wd = pattern(320 * 10240, 3).reshape(320, 10240)
        wu = pattern(10240 * 320, 4).reshape(10240, 320)
        wi = pattern(4 * 10240, 5).reshape(4, 10240)
        emit(f"hyper_inj_s{rows}", fused_hyper_read(x, gamma, wd, wu, wi))
        mixed, _ = fused_hyper_read(x, gamma, wd, wu)
        emit(f"hyper_mix_s{rows}", [mixed])
    # Use the reference's own quantizer and module packing for v3.
    from types import SimpleNamespace
    from mtplx.kernels.hyper_connection_v3 import prepare_v3_pack
    pack = prepare_v3_pack(SimpleNamespace(
        input_mix_weight_down=SimpleNamespace(weight=wd),
        input_mix_weight_up=SimpleNamespace(weight=wu),
        block_inject_weight=SimpleNamespace(weight=wi)))
    emit("hyper_v3", fused_hyper_read_v3(pattern(10240, 1), gamma, pack))
    for rows in range(1, 7):
        x = pattern(rows * 10240, 6).reshape(rows, 10240)
        state = pattern(3 * 10240, 7).reshape(3, 10240)
        weights = pattern(10240 * 4, 8).reshape(10240, 4)
        outputs = (fused_gdn_conv_norm(x.reshape(-1), state, weights) if rows == 1
                   else fused_gdn_conv_norm_rows(x, state, weights))
        emit(f"gdn_s{rows}", outputs)
    # Carry both returned states, including nonzero FP32 recurrence, across steps.
    delta = pattern(48 * 128 * 128, 15).astype(mx.float32).reshape(48, 128, 128)
    z, a, b = pattern(6144, 9), pattern(48, 10), pattern(48, 11)
    a_log, dt_bias, norm = pattern(48, 12), pattern(48, 13), pattern(128, 14)
    for step in range(2):
        y, state, delta = fused_gdn_step(
            pattern(10240, 6 + step), z, a, b, state, weights,
            a_log, dt_bias, delta, norm)
        emit(f"gdn_step_{step}", [y, state, delta])
    # Synthetic valid affine packs isolate this kernel's ABI from the quantizer.
    words = (np.arange(2560 * 768, dtype=np.uint64) * 2654435761 + 12345).astype(np.uint32)
    qw = mx.array(words).reshape(2560, 768)
    for gs in (32, 64):
        scales = pattern(2560 * 6144 // gs, 16).reshape(2560, 6144 // gs)
        biases = pattern(2560 * 6144 // gs, 17).reshape(2560, 6144 // gs)
        for tag, dtype in (("bf16", mx.bfloat16), ("f32", mx.float32)):
            x = pattern(6144, 18).astype(dtype)
            emit(f"gdn_out_gs{gs}_{tag}", [fused_gdn_out(
                x, z, norm, qw, scales, biases, group_size=gs)])
    for rows in (1, 4, 7, 32, 2048):
        q = (pattern(rows * 2048, 19) / 16).reshape(1, rows, 16, 128)
        k = (pattern(rows * 2048, 20) / 16).reshape(1, rows, 16, 128)
        v = pattern(rows * 6144, 21).reshape(1, rows, 48, 128)
        g = mx.full((1, rows, 48), 0.75, dtype=mx.float32)
        beta = mx.full((1, rows, 48), 0.5, dtype=mx.bfloat16)
        state = pattern(48 * 128 * 128, 15).astype(mx.float32).reshape(1, 48, 128, 128)
        for masked in (False, True):
            mask = mx.array((np.arange(rows) % 3) != 0).reshape(1, rows) if masked else None
            emit(f"gated_delta_s{rows}_mask{int(masked)}", gated_delta.gated_delta_kernel(
                q, k, v, g, beta, state, mask))
    # Exercise real runtime selection, not individually chosen fast kernels.
    shapes = [(1, 80, 64), (4, 80, 128), (1, 2560, 2560), (1, 324, 320),
              (1, 81, 2560), *[(m, 640, 2560) for m in (2, 3, 4, 5, 6, 12, 24, 25)],
              (13, 6144, 6144), (33, 48, 2560), (13, 16384, 320),
              (65, 16385, 320), (13, 16384, 160), (2048, 2560, 2560),
              (33, 16, 32768)]
    for m, n, k in shapes:
        for group in (32, 64):
            if k % group:
                continue
            for bits in (4, 8):
                words = (np.arange(n * k * bits // 32, dtype=np.uint64)
                         * 2654435761 + 12345).astype(np.uint32)
                qw = mx.array(words).reshape(n, k * bits // 32)
                qs = pattern(n * k // group, 23).reshape(n, k // group)
                qb = pattern(n * k // group, 24).reshape(n, k // group)
                x = pattern(m * k, 25).reshape(m, k)
                y = mx.quantized_matmul(x, qw, qs, qb, transpose=True, group_size=group, bits=bits)
                emit(f"affine_m{m}_n{n}_k{k}_g{group}_b{bits}", [y],
                     m=m, n=n, k=k, group=group, bits=bits)
    # Call the actual model's live-state branch, not a hand-written stand-in.
    import mlx.nn as nn
    from mlx_lm.models.cache import ArraysCache
    from mtplx.models.qwen4_exp import GatedDeltaNet, TextArgs, _FusedGDNInProj
    os.environ["MTPLX_FUSED_GDN_STEP"] = "1"
    layer = GatedDeltaNet(TextArgs())
    def pack(n, k, salt):
        words = (np.arange(n * k // 8, dtype=np.uint64) * 2654435761 + 12345).astype(np.uint32)
        return (mx.array(words).reshape(n,k//8),
                (pattern(n*k//64,salt)/64).reshape(n,k//64),
                (pattern(n*k//64,salt+1)/64).reshape(n,k//64))
    layer.in_proj_fused = _FusedGDNInProj(*pack(16480,2560,26),64,4,"affine",[10240,16384,16432])
    layer.out_proj = nn.QuantizedLinear(6144,2560,bias=False,group_size=64,bits=4)
    layer.out_proj.weight, layer.out_proj.scales, layer.out_proj.biases = pack(2560,6144,28)
    layer.conv1d.weight = pattern(10240*4,8).reshape(10240,4,1)
    layer.A_log, layer.dt_bias, layer.norm.weight = pattern(48,12), pattern(48,13), pattern(128,14)
    layer.eval()
    cache = ArraysCache(2)
    cache[0] = pattern(3*10240,7).reshape(1,3,10240)
    cache[1] = pattern(48*128*128,15).astype(mx.float32).reshape(1,48,128,128)
    for step in range(2):
        assert layer._fused_step_applies(1,1,None,cache)
        output = layer(pattern(2560,30+step).reshape(1,1,2560),cache=cache)
        emit(f"gdn_chain_{step}",[output,cache[0],cache[1]])
    # Full staged model forward, including captured recurrence inputs and
    # carried states. Keep actual module selection and BF16 boundaries.
    from mtplx.models.qwen4_exp import _VERIFY_CAPTURE
    os.environ["MTPLX_FUSED_GDN_STEP"] = "0"
    staged_cases = [(rows,initial,masked,fused,None)
                    for rows in (1,2,4,6,7,32,2048)
                    for initial in (False,True) for masked in (False,True)
                    for fused in (False,True)]
    staged_cases += [(rows,True,False,True,length)
                     for rows in (1,4,7) for length in (-3,rows-1,rows+3)]
    for rows,initial,masked,fused,length in staged_cases:
        for flag in ("MTPLX_FUSED_GDN_CONVNORM","MTPLX_FUSED_CONVNORM_VERIFY","MTPLX_FUSED_GDN_OUT"):
            os.environ[flag] = str(int(fused))
        cache = ArraysCache(2)
        if initial:
            cache[0] = pattern(3*10240,7).reshape(1,3,10240)
            cache[1] = pattern(48*128*128,15).astype(mx.float32).reshape(1,48,128,128)
        mask = mx.array((np.arange(rows)%3)!=0).reshape(1,rows) if masked else None
        for step in range(2):
            # The host owns cache lengths; exercise clipping at either end.
            cache.lengths = mx.array([length],dtype=mx.int32) if length is not None else None
            used_fused = layer._fused_conv_norm_applies(1,rows,mask,cache) or layer._fused_conv_norm_rows_applies(1,rows,mask,cache)
            capture = _VERIFY_CAPTURE.set(True)
            try:
                output = layer(pattern(rows*2560,30+step).reshape(1,rows,2560),mask=mask,cache=cache)
            finally:
                _VERIFY_CAPTURE.reset(capture)
            _,q,k,v,_,_ = cache._mtplx_verify_rows
            emit(f"gdn_staged_r{rows}_initial{int(initial)}_mask{int(masked)}_fused{int(fused)}_len{length}_step{step}",
                 [output,cache[0],cache[1],q,k,v],rows=rows,initial=initial,masked=masked,fused=fused,
                 length=length,step=step,used_fused=used_fused)
    # Unified model branch selection, including non-sanitized input modules.
    forward_cases = [(1,initial,mode) for initial in range(4) for mode in range(6)]
    forward_cases += [(rows,initial,mode) for rows in (2,6,7,33,2048)
                      for initial in (0,1) for mode in (0,1)]
    for layout in ("fused","separate","mixed"):
        if layout != "fused":
            if hasattr(layer,"in_proj_fused"):
                del layer.in_proj_fused
            for index,(name,width) in enumerate(zip(("qkv","z","b","a"),(10240,6144,48,48))):
                bits,group = ((4,64),(8,32),(8,64),(4,32))[index] if layout=="mixed" else (4,64)
                proj = nn.QuantizedLinear(2560,width,bias=False,group_size=group,bits=bits)
                words = (np.arange(width*2560*bits//32,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
                proj.weight = mx.array(words).reshape(width,2560*bits//32)
                proj.scales = (pattern(width*2560//group,64+index*2)/64).reshape(width,2560//group)
                proj.biases = (pattern(width*2560//group,65+index*2)/64).reshape(width,2560//group)
                setattr(layer,"in_proj_"+name,proj)
        for rows,initial,mode in forward_cases:
            enabled = mode != 0
            for flag in ("MTPLX_FUSED_GDN_STEP","MTPLX_FUSED_GDN_CONVNORM","MTPLX_FUSED_CONVNORM_VERIFY","MTPLX_FUSED_GDN_OUT"):
                os.environ[flag] = str(int(enabled and (mode!=5 or flag=="MTPLX_FUSED_GDN_STEP")))
            cache = ArraysCache(2)
            if initial in (1,3): cache[0] = pattern(3*10240,7).reshape(1,3,10240)
            if initial in (1,2): cache[1] = pattern(48*128*128,15).astype(mx.float32).reshape(1,48,128,128)
            if mode == 4: cache.lengths = mx.array([1],dtype=mx.int32)
            mask = mx.array((np.arange(rows)%3)!=0).reshape(1,rows) if mode==3 else None
            for step in range(2):
                capture = _VERIFY_CAPTURE.set(mode==2)
                try:
                    used_step = layer._fused_step_applies(1,rows,mask,cache)
                    used_conv = not used_step and (layer._fused_conv_norm_applies(1,rows,mask,cache) or layer._fused_conv_norm_rows_applies(1,rows,mask,cache))
                    output = layer(pattern(rows*2560,30+step).reshape(1,rows,2560),mask=mask,cache=cache)
                finally:
                    _VERIFY_CAPTURE.reset(capture)
                emit(f"gdn_forward_{layout}_r{rows}_initial{initial}_mode{mode}_step{step}",
                     [output,cache[0],cache[1]],layout=layout,rows=rows,initial=initial,mode=mode,
                     step=step,used_step=used_step,used_conv=used_conv)
    from mtplx.models.qwen4_exp import QSAIndexer
    from mtplx.kernels.qsa_indexer_prepare import qsa_indexer_prepare_queries_metal, qsa_indexer_pool_keys_metal
    indexer = QSAIndexer(TextArgs())
    indexer.q_layernorm.weight = pattern(128,72)
    indexer.k_layernorm.weight = pattern(128,73)
    frequencies = np.asarray(indexer._inv_freq).tolist()
    assert len(frequencies)==32 and indexer._rope_attention_scaling==1.0
    for rows in (1,4,7,32,2048):
        for pos in (0,32767,262141):
            for stride,dimstride in ((512,1),(640,1),(1280,2)):
                raw = pattern(rows*stride,71).reshape(1,rows,stride)[...,:512*dimstride:dimstride].reshape(1,rows,4,128)
                q = qsa_indexer_prepare_queries_metal(raw,indexer.q_layernorm.weight,indexer._inv_freq,
                    pos_start=pos,eps=1e-6)
                emit(f"qsa_prepare_r{rows}_stride{stride}_dim{dimstride}_pos{pos}",[q],
                    rows=rows,stride=stride,dimstride=dimstride,pos=pos,frequencies=frequencies)
                q = indexer._prepare_queries_eager(raw,pos)
                emit(f"qsa_prepare_eager_r{rows}_stride{stride}_dim{dimstride}_pos{pos}",[q],
                    rows=rows,stride=stride,dimstride=dimstride,pos=pos,frequencies=frequencies,eager=True)
    for blocks in (1,2,7,512):
        for start in (0,8191,65535):
            for stride,dimstride in ((128,1),(640,1),(1280,2)):
                raw = pattern(blocks*4*stride,74).reshape(1,blocks*4,stride)[...,:128*dimstride:dimstride]
                pooled = qsa_indexer_pool_keys_metal(raw,indexer.k_layernorm.weight,indexer._inv_freq,
                    block_start=start,compress_ratio=4,eps=1e-6)
                emit(f"qsa_pool_n{blocks}_stride{stride}_dim{dimstride}_start{start}",[pooled],
                    blocks=blocks,stride=stride,dimstride=dimstride,start=start,frequencies=frequencies)
                pooled = indexer._pool_keys_eager(raw,start,start+blocks)
                emit(f"qsa_pool_eager_n{blocks}_stride{stride}_dim{dimstride}_start{start}",[pooled],
                    blocks=blocks,stride=stride,dimstride=dimstride,start=start,frequencies=frequencies,eager=True)
    os.environ["MTPLX_FUSED_QSA_INDEXER"] = "1"
    for bits,group in ((4,64),(8,64),(8,32)):
        proj = nn.QuantizedLinear(2560,640,bias=False,group_size=group,bits=bits)
        words = (np.arange(640*2560*bits//32,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
        proj.weight = mx.array(words).reshape(640,2560*bits//32)
        proj.scales = (pattern(640*2560//group,75)/64).reshape(640,2560//group)
        proj.biases = (pattern(640*2560//group,76)/64).reshape(640,2560//group)
        indexer.index_qk_proj = proj
        for rows in (4,8,32,2048):
            for pos in (0,262140):
                qk = indexer.index_qk_proj(pattern(rows*2560,77).reshape(1,rows,2560))
                q,k = mx.split(qk,[512],axis=-1)
                prepared = indexer._prepare_queries(q.reshape(1,rows,4,128),pos)
                pooled = qsa_indexer_pool_keys_metal(k,indexer.k_layernorm.weight,indexer._inv_freq,
                    block_start=pos//4,compress_ratio=4,eps=1e-6)
                emit(f"qsa_project_r{rows}_b{bits}_g{group}_pos{pos}",[prepared,pooled],
                    rows=rows,bits=bits,group=group,pos=pos,frequencies=frequencies)
                prepared = indexer._prepare_queries_eager(q.reshape(1,rows,4,128),pos)
                pooled = indexer._pool_keys_eager(k,pos//4,pos//4+rows//4)
                emit(f"qsa_project_eager_r{rows}_b{bits}_g{group}_pos{pos}",[prepared,pooled],
                    rows=rows,bits=bits,group=group,pos=pos,frequencies=frequencies,eager=True)
    # Actual QSACache + KVCache lifecycle, retaining array aliases across restore.
    from mtplx.models.qwen4_exp import QSACache
    cache_ops = [
        ("append",1),("append",3),("append",252),("save",0),("append",1),("append",767),
        ("save",1),("append",4),("trim",3),("append",7),("restore",1),("append",4),
        ("restore",0),("trim",5),("append",2),("reserve",(1600,600)),("append",1029),
        ("save",2),("trim",100000),("append",2),("restore",2),("append",3),
    ]
    for scenario in (0,1,2,3):
        os.environ["MTPLX_FUSED_QSA_INDEXER"] = "1" if scenario < 2 else "0"
        cache = QSACache()
        bank = {}
        if scenario%2: cache.reserve_indexer_capacity(raw_capacity=300,pooled_capacity=300)
        for index,(op,arg) in enumerate(cache_ops):
            if op=="append":
                stride=128 if index%2==0 else 640
                raw=pattern(arg*stride,80+index).reshape(1,arg,stride)[...,:128]
                keys=pattern(arg*512,110+index).reshape(1,arg,2,256).transpose(0,2,1,3)
                values=pattern(arg*512,140+index).reshape(1,arg,2,256).transpose(0,2,1,3)
                cache.write_raw(raw)
                indexer._extend_pooled(cache,cache.offset+arg)
                cache.kv.update_and_fetch(keys,values)
            elif op=="reserve": cache.reserve_indexer_capacity(raw_capacity=arg[0],pooled_capacity=arg[1])
            elif op=="trim": cache.trim(arg)
            elif op=="save": bank[arg]=cache.state
            elif op=="restore": cache.state=bank[arg]
            # A mirror read after reserve/restore must seed every earlier block.
            if cache.pooled_len: cache.pooled_f32_view(cache.pooled_len)
            arrays=[cache.kv.keys,cache.kv.values,cache.raw_keys,cache.pooled,cache.pooled_f32_t]
            outputs=[x for x in arrays if x is not None]
            for saved in bank.values(): outputs.extend(x for x in saved if x is not None)
            emit(f"qsa_cache_s{scenario}_op{index}",outputs,scenario=scenario,index=index,op=op,arg=arg,
                 offset=cache.offset,nbytes=cache.nbytes,raw_cap=cache.raw_keys.shape[1],kv_cap=cache.kv.keys.shape[2],
                 pooled_cap=0 if cache.pooled is None else cache.pooled.shape[1],pooled_len=cache.pooled_len,
                 mirror_cap=0 if cache.pooled_f32_t is None else cache.pooled_f32_t.shape[3],
                 present=[x is not None for x in arrays],bank=list(bank),frequencies=frequencies)
    os.environ["MTPLX_FUSED_QSA_INDEXER"] = "1"
    # Original fused QSA selector: all output modes, padded histories, tie order,
    # BF16/F32 operands and the runtime TF32 setting. Also capture its actual
    # header/body, independent of our Rust entry-point binding declaration.
    from mtplx.kernels import qsa_indexer_select as selector
    select_shapes = [(1,0,1,0,512),(4,0,7,8,512),(4,2045,2049,600,512),
                     (7,2097,2104,600,16),(33,4080,4113,1200,512),(1,8192,8193,4096,512)]
    for case,(rows,pos,total,blocks,topk) in enumerate(select_shapes):
        for variant in range(4):
            f32 = [variant==2, variant>=2]
            tf32 = variant!=3
            dim_stride = 2 if variant%2 else 1
            qstride=512*dim_stride
            pstride=128*dim_stride
            def operand(count,salt,fp32):
                x=pattern(max(1,count),salt)
                if fp32: x=x.astype(mx.float32)+0.00013
                if variant==1: x=mx.zeros_like(x)
                return x
            q=operand(rows*qstride,90,f32[0]).reshape(1,rows,4,128*dim_stride)[...,::dim_stride]
            pooled=operand(blocks*pstride,91,f32[1])[:blocks*pstride].reshape(1,blocks,128*dim_stride)[...,::dim_stride]
            for mode in ("blocks","dense_mask","row_tokens"):
                dense=(blocks+1)*4 if mode=="dense_mask" else 0
                metal_kernel=mx.fast.metal_kernel
                try:
                    mx.fast.metal_kernel=lambda **kw:kw
                    spec=selector._selector_kernel.__wrapped__(mode,4,128,blocks,topk,4,
                        1<<(max(256,topk)-1).bit_length(),
                        dense,tf32,q.dtype,pooled.dtype)
                finally: mx.fast.metal_kernel=metal_kernel
                tf_setting=selector._mlx_tf32_enabled
                try:
                    selector._mlx_tf32_enabled=lambda:tf32
                    out=selector.qsa_indexer_select_metal(q,pooled,pos_start=pos,total_tokens=total,
                        logical_blocks=total//4,block_topk=topk,compress_ratio=4,mode=mode,
                        output_total_tokens=dense if mode=="dense_mask" else None)
                finally: selector._mlx_tf32_enabled=tf_setting
                emit(f"qsa_select_c{case}_v{variant}_{mode}",[out] if mode=="dense_mask" else out,
                    rows=rows,pos=pos,total=total,blocks=blocks,topk=topk,mode=mode,dense_width=dense,
                    input_f32=f32,tf32=tf32,dim_stride=dim_stride,zero=variant==1,
                    header_sha256=hashlib.sha256(spec["header"].encode()).hexdigest(),
                    body_sha256=hashlib.sha256(spec["source"].encode()).hexdigest())
    # Force the actual 32MiB scratch boundary with a large reserved backing
    # but a short logical prefix. This calls the host chunk/concat path itself.
    q=pattern(129*512,90).reshape(1,129,4,128)
    pooled=pattern(65536*128,91).reshape(1,65536,128)
    for mode in ("blocks","dense_mask","row_tokens"):
        out=indexer._select_fused(q,0,pooled,129//4,129,mode)
        emit(f"qsa_chunk_{mode}",[out] if mode=="dense_mask" else out,
             rows=129,pos=0,total=129,blocks=65536,topk=512,mode=mode,dense_width=0,
             input_f32=[False,False],tf32=True,dim_stride=1,zero=False,chunked=True)
    # Connected raw-write -> pool -> prepare -> selection -> KV advance.
    cache=QSACache()
    cache.write_raw(pattern(2048*128,88).reshape(1,2048,128))
    indexer._extend_pooled(cache,2048)
    kv=pattern(2048*512,89).reshape(1,2048,2,256).transpose(0,2,1,3)
    cache.kv.update_and_fetch(kv,kv)
    for step,rows in enumerate((1,4,7,2048)):
        q=indexer._prepare_queries(pattern(rows*512,90+step).reshape(1,rows,4,128),cache.offset)
        cache.write_raw(pattern(rows*128,100+step).reshape(1,rows,128))
        total=cache.offset+rows
        indexer._extend_pooled(cache,total)
        for mode in ("blocks","dense_mask","row_tokens"):
            if rows>=128:
                from mtplx.kernels.qsa_indexer_prefill import qsa_indexer_prefill_metal
                out=qsa_indexer_prefill_metal(q,cache.pooled[:,:total//4,:],pos_start=cache.offset,
                    total_tokens=total,block_topk=512,compress_ratio=4,logical_blocks=total//4,mode=mode)
            else: out=indexer._select_fused(q,cache.offset,cache.pooled,total//4,total,mode)
            emit(f"qsa_connected_s{step}_{mode}",[out] if mode=="dense_mask" else out,
                 rows=rows,step=step,mode=mode,frequencies=frequencies,**({"prefill":True} if rows>=128 else {}))
        kv=pattern(rows*512,110+step).reshape(1,rows,2,256).transpose(0,2,1,3)
        cache.kv.update_and_fetch(kv,kv)
    from mtplx.kernels import qsa_indexer_prefill as prefill
    for rows,blocks in ((1,1),(4,7),(15,31),(16,32),(17,33),(33,513),(2048,513)):
        for ds in (1,2):
            q=pattern(rows*512*ds,92).reshape(1,rows,4,128*ds)[...,::ds]
            pooled=pattern(blocks*128*ds,93).reshape(1,blocks,128*ds)[...,::ds]
            assert prefill.qsa_indexer_prefill_scores_mpp_supported(q,pooled)
            emit(f"qsa_mpp_score_r{rows}_n{blocks}_ds{ds}",[prefill.qsa_indexer_prefill_scores_mpp(q,pooled)],
                 rows=rows,blocks=blocks,dim_stride=ds)
    topk_shapes=((1,0,1,1,512),(4,0,7,8,512),(7,2045,2052,600,512),
                 (4,8192,8196,4096,512),(1,262142,262143,65536,512),
                 (4,252,256,64,16),(1,256,257,65,16),(1,8188,8189,2048,512),(1,8192,8193,2049,512))
    for case,(rows,pos,total,blocks,topk) in enumerate(topk_shapes):
        for variant in range(4):
            ds=2 if variant in (1,3) else 1
            stride=blocks*ds+(4 if variant else 0)
            scores=pattern(rows*stride,94).astype(mx.float32)
            if variant in (1,2): scores=mx.full((rows*stride,),float(variant-1),dtype=mx.float32)
            scores=scores.reshape(rows,stride)[:,:blocks*ds:ds]
            for mode in ("blocks","dense_mask","row_tokens"):
                dense=(blocks+1)*4 if mode=="dense_mask" else 0
                metal_kernel=mx.fast.metal_kernel
                try:
                    mx.fast.metal_kernel=lambda **kw:kw
                    spec=prefill._prefill_topk_kernel.__wrapped__(mode,blocks,topk,4,
                        1<<(max(256,topk)-1).bit_length(),dense)
                finally: mx.fast.metal_kernel=metal_kernel
                out=prefill.qsa_indexer_prefill_topk_metal(scores,pos_start=pos,total_tokens=total,
                    block_topk=topk,compress_ratio=4,logical_blocks=total//4,mode=mode,
                    output_total_tokens=dense if mode=="dense_mask" else None)
                emit(f"qsa_prefill_topk_c{case}_v{variant}_{mode}",[out] if mode=="dense_mask" else out,
                     rows=rows,pos=pos,total=total,blocks=blocks,topk=topk,mode=mode,dense_width=dense,
                     variant=variant,row_stride=stride,dim_stride=ds,
                     header_sha256=hashlib.sha256(spec["header"].encode()).hexdigest(),
                     body_sha256=hashlib.sha256(spec["source"].encode()).hexdigest())
    # Producer-aware budget/planner, and the actual score -> topk -> concat graph.
    for rows,blocks,budget in ((7,513,513*4*3),(65,600,600*4*48),(2048,513,128*1024*1024),
                              (3,65536,1),(519,65536,128*1024*1024)):
        for mpp in (False,True):
            chunk=prefill.qsa_indexer_prefill_score_chunk_rows(rows,4,blocks,budget,producer="mpp" if mpp else "mlx")
            emit(f"qsa_prefill_plan_r{rows}_n{blocks}_b{budget}_mpp{int(mpp)}",[],rows=rows,blocks=blocks,budget=budget,mpp=mpp,chunk=chunk)
    for case,(rows,blocks,budget) in enumerate(((7,513,513*4*3),(65,600,600*4*48),(2048,513,128*1024*1024),(3,65536,1))):
        for ds in (1,2):
            q=pattern(rows*512*ds,92).reshape(1,rows,4,128*ds)[...,::ds]
            pooled=pattern(blocks*128*ds,93).reshape(1,blocks,128*ds)[...,::ds]
            total=max(rows,blocks*4-1)
            pos=total-rows
            for mode in ("blocks","dense_mask","row_tokens"):
                dense=total+4 if mode=="dense_mask" else 0
                out=prefill.qsa_indexer_prefill_metal(q,pooled,pos_start=pos,total_tokens=total,
                    block_topk=512,compress_ratio=4,logical_blocks=total//4,mode=mode,
                    output_total_tokens=dense if mode=="dense_mask" else None,score_workspace_bytes=budget)
                emit(f"qsa_prefill_chain_c{case}_ds{ds}_{mode}",[out] if mode=="dense_mask" else out,
                     rows=rows,pos=pos,total=total,blocks=blocks,topk=512,mode=mode,dense_width=dense,
                     budget=budget,dim_stride=ds)
    # Actual gather selection, including the distinct sorted-RHS threshold.
    gather_shapes = [
        (3, 1, 640, 2560, 8, 40, False),
        (3, 1, 2560, 640, 8, 40, False),
        (3, 1, 81, 320, 8, 7, False),
        (3, 33, 80, 128, 8, 7, False),
        (3, 33, 81, 160, 8, 7, False),
        (31, 1, 80, 128, 8, 31, True),
        (32, 1, 80, 128, 8, 32, True),
        (33, 1, 81, 160, 8, 33, True),
        (512, 1, 1280, 2560, 8, 512, True),
        (2047, 1, 80, 128, 512, 2047, True),
        (2048, 1, 80, 128, 512, 2048, True),
    ]
    for batches, m, n, k, experts, routes, sorted_indices in gather_shapes:
        for group in (32, 64):
            if k % group:
                continue
            for bits in (4, 8):
                words = (np.arange(experts*n*k*bits//32, dtype=np.uint64)
                         * 2654435761 + 12345).astype(np.uint32)
                def padded_view(value, shape, extra):
                    storage = mx.concatenate([value,mx.zeros((extra,),dtype=value.dtype)])
                    mx.eval(storage)
                    return storage[:value.size].reshape(shape)
                padding = k if sorted_indices and k % 64 else 0
                qw = padded_view(mx.array(words),(experts,n,k*bits//32),padding*bits//32*64)
                qs = padded_view(pattern(experts*n*k//group,33),(experts,n,k//group),padding//group*64)
                qb = padded_view(pattern(experts*n*k//group,34),(experts,n,k//group),padding//group*64)
                # The pinned RHS NAX shader reads a full 64-wide final tile
                # even for K=160. Keep identical physical padding in both
                # runners; this synthetic shape is not an installed MoE shape.
                x = padded_view(pattern(batches*m*k,32),(batches,m,k),padding*64)
                rhs = (np.arange(routes,dtype=np.uint32)*5+2)%experts
                if sorted_indices:
                    rhs.sort()
                lhs = None if sorted_indices else mx.array(
                    (np.arange(routes,dtype=np.uint32)*2+1)%batches)
                y = mx.gather_qmm(x,qw,qs,qb,lhs_indices=lhs,rhs_indices=mx.array(rhs),
                                  transpose=True,group_size=group,bits=bits,
                                  sorted_indices=sorted_indices)
                emit(f"gather_x{batches}_m{m}_n{n}_k{k}_e{experts}_r{routes}_s{int(sorted_indices)}_g{group}_b{bits}",
                     [y],batches=batches,m=m,n=n,k=k,experts=experts,routes=routes,
                     sorted=sorted_indices,group=group,bits=bits,padding=padding)
    # Routing uses complete last-axis sort; gather_sort sorts the returned order
    # again for its inverse. Check tied values, tile boundaries and long merges.
    shapes = [(rows,width) for rows in (1,3) for width in (
        1,10,63,64,65,128,129,256,257,512,513,1024,1025,2048,2049,
        4097,20480,65535,65536,65537)] + [(2048,512)]
    for rows,width in shapes:
        for tag,dtype in (("u32",mx.uint32),("bf16",mx.bfloat16),("f32",mx.float32)):
            values = (np.arange(rows*width,dtype=np.uint32)*17+35)%512
            x = mx.array(values,dtype=dtype).reshape(rows,width)
            order = mx.argsort(x,axis=-1)
            inverse = mx.argsort(order,axis=-1)
            emit(f"argsort_{tag}_r{rows}_w{width}",[order,inverse],rows=rows,width=width,dtype=tag)
    from mlx_lm.models.switch_layers import _gather_sort, _scatter_unsort
    for tokens,top_k,width in ((1,10,7),(6,10,8),(7,10,9),(7,10,2560),
                               (205,10,9),(2048,10,2560),(6554,10,9)):
        x = pattern(tokens*width,36).reshape(1,tokens,1,1,width)
        indices = mx.array((np.arange(tokens*top_k,dtype=np.uint32)*17+35)%512).reshape(1,tokens,top_k)
        sorted_x, sorted_indices, inverse = _gather_sort(x,indices)
        unsorted = _scatter_unsort(sorted_x,inverse,indices.shape)
        emit(f"gather_sort_t{tokens}_top{top_k}_w{width}",
             [sorted_x,sorted_indices,inverse,unsorted],tokens=tokens,top_k=top_k,width=width)
    for rows,intermediate in ((1,64),(1,640),(1,65534),(1,65536),(1,65538),
                              (2,640),(10,640),(60,640),(70,640),(20480,640)):
        x = pattern(rows*intermediate*2,37).reshape(rows,intermediate*2)
        gate,up = mx.split(x,2,axis=-1)
        activated = nn.silu(gate)
        emit(f"gate_up_r{rows}_i{intermediate}",[activated,activated*up],rows=rows,intermediate=intermediate)
    # Exact SparseMoeBlock routing operations, retaining the strided top-k view.
    for rows,experts in ((1,512),(7,512),(32,512),(2048,512),
                         (1,16),(7,16),(3,4096),(3,4097)):
        for normalize in (False,True):
            logits = pattern(rows*experts,43).reshape(rows,experts)
            gates = mx.softmax(logits,axis=-1,precise=True)
            indices = mx.argpartition(gates,kth=-10,axis=-1)[...,-10:]
            scores = mx.take_along_axis(gates,indices,axis=-1)
            if normalize:
                scores = scores/scores.sum(axis=-1,keepdims=True)
            emit(f"router_r{rows}_e{experts}_norm{int(normalize)}",
                 [gates,indices,scores],rows=rows,experts=experts,normalize=normalize)
    from mlx_lm.models.switch_layers import QuantizedSwitchLinear
    from mtplx.models.qwen4_exp import _FusedGateUpSwitchGLU
    def expert_pack(experts,n,k,group,salt):
        words = (np.arange(experts*n*k//8,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
        return (mx.array(words).reshape(experts,n,k//8),
                (pattern(experts*n*k//group,salt)/64).reshape(experts,n,k//group),
                (pattern(experts*n*k//group,salt+1)/64).reshape(experts,n,k//group))
    for tokens,experts,k,intermediate in (
        (1,16,2560,640),(6,16,2560,640),(7,16,2560,640),
        (205,16,2560,640),(2048,16,2560,640),(7,512,128,64),(205,512,128,64)):
        for gu_group,down_group in ((32,32),(64,64),(32,64)):
            # The reference module properties derive dimensions from loaded weights.
            # Small constructor placeholders avoid unrelated random full-pack work.
            down = QuantizedSwitchLinear(64,64,1,bias=False,group_size=down_group,bits=4)
            down.weight,down.scales,down.biases = expert_pack(experts,k,intermediate,down_group,41)
            layer = _FusedGateUpSwitchGLU(down,*expert_pack(experts,2*intermediate,k,gu_group,39),gu_group,4,"affine")
            layer.eval()
            x = pattern(tokens*k,38).reshape(1,tokens,k)
            idx = mx.array((np.arange(tokens*10,dtype=np.uint32)*17+35)%experts).reshape(1,tokens,10)
            emit(f"switch_t{tokens}_e{experts}_k{k}_i{intermediate}_g{gu_group}_d{down_group}",
                 [layer(x,idx)],tokens=tokens,experts=experts,k=k,intermediate=intermediate,
                 gu_group=gu_group,down_group=down_group)
    # Call the complete actual stock SparseMoeBlock, including router projection,
    # expert weighting and shared expert. Opt-in fused branches get separate receipts.
    from mtplx.models.qwen4_exp import SparseMoeBlock, _FusedGateUpMLP
    from mlx_lm.models.switch_layers import SwitchGLU
    from mlx_lm.models.qwen3_next import Qwen3NextMLP
    os.environ["MTPLX_FUSED_MOE_DECODE"] = "0"
    os.environ["MTPLX_FUSED_MOE_VERIFY"] = "0"
    linear = quantized_linear_fixture
    moe_shapes = (
        (1,16,2560,640),(4,16,2560,640),(6,16,2560,640),(7,16,2560,640),
        (25,16,2560,640),(205,16,2560,640),(2048,16,2560,640),
        (7,512,128,64),(205,512,128,64),
        (1,512,2560,640),(7,512,2560,640),(2048,512,2560,640))
    moe_cases = [(shape,((64,64),(32,64)),False,(False,False)) for shape in moe_shapes]
    moe_cases += [((tokens,experts,2560,640),((32,32),(32,64),(64,32),(64,64)),True,(False,False))
                  for tokens,experts in ((1,16),(2,16),(3,16),(4,16),(5,16),(7,16),(1,512),(4,512))]
    moe_cases += [((tokens,experts,2560,640),((64,64),(32,64)),True,separate)
                  for tokens,experts in ((1,16),(4,16),(6,16),(7,16),(25,16),(2048,16),(1,512),(7,512),(2048,512))
                  for separate in ((True,False),(False,True),(True,True))]
    for (tokens,experts,k,intermediate),groups,fused,(separate_experts,separate_shared) in moe_cases:
        os.environ["MTPLX_FUSED_MOE_DECODE"] = str(int(fused))
        os.environ["MTPLX_FUSED_MOE_VERIFY"] = str(int(fused))
        for gu_group,down_group in groups:
            layer = SparseMoeBlock(SimpleNamespace(hidden_size=64,moe_intermediate_size=64,
                shared_expert_intermediate_size=64,norm_topk_prob=True,num_experts=1,num_experts_per_tok=10))
            down = QuantizedSwitchLinear(64,64,1,bias=False,group_size=down_group,bits=4)
            down.weight,down.scales,down.biases = expert_pack(experts,k,intermediate,down_group,41)
            if separate_experts:
                sw = SwitchGLU(64,64,1)
                for name,group,salt in (("gate_proj",gu_group,39),("up_proj",down_group,45)):
                    proj = QuantizedSwitchLinear(64,64,1,bias=False,group_size=group,bits=4)
                    proj.weight,proj.scales,proj.biases = expert_pack(experts,intermediate,k,group,salt)
                    setattr(sw,name,proj)
                sw.down_proj = down
                layer.switch_mlp = sw
            else:
                layer.switch_mlp = _FusedGateUpSwitchGLU(down,*expert_pack(experts,2*intermediate,k,gu_group,39),gu_group,4,"affine")
            if separate_shared:
                shared = Qwen3NextMLP(64,64)
                shared.gate_proj = linear(intermediate,k,4,gu_group,46)
                shared.up_proj = linear(intermediate,k,4,down_group,54)
                shared.down_proj = linear(k,intermediate,4,down_group,48)
                layer.shared_expert = shared
            else:
                shared_gu = linear(2*intermediate,k,4,gu_group,46)
                layer.shared_expert = _FusedGateUpMLP(linear(k,intermediate,4,down_group,48),
                    shared_gu.weight,shared_gu.scales,shared_gu.biases,gu_group,4,"affine")
            layer.gate = linear(experts,k,8,64,50)
            layer.shared_expert_gate = linear(1,k,4,64,52)
            layer.num_experts = experts
            layer.eval()
            x = pattern(tokens*k,44).reshape(1,tokens,k)
            for normalize in (False,True):
                layer.norm_topk_prob = normalize
                extra = {"fused":True} if fused else {}
                tag = "moe_fused" if fused else "moe"
                if separate_experts or separate_shared:
                    extra.update(separate_experts=separate_experts,separate_shared=separate_shared)
                    tag = f"moe_separate_e{int(separate_experts)}_s{int(separate_shared)}"
                emit(f"{tag}_t{tokens}_e{experts}_k{k}_i{intermediate}_g{gu_group}_d{down_group}_norm{int(normalize)}",
                     [layer(x)],tokens=tokens,experts=experts,k=k,intermediate=intermediate,
                     gu_group=gu_group,down_group=down_group,normalize=normalize,**extra)
    from mtplx.models.qwen4_exp import GroupedRMSNorm, SigmoidRMSNormGated
    for rows,width,group in ((1,10240,2560),(4,10240,2560),(2048,10240,2560),
                             (1,65534,2),(1,65536,128),(1,65538,2),(3,32768,8192),
                             (1,256,128),(7,128,128),(3,16384,4096)):
        norm = GroupedRMSNorm(width,group,eps=1e-6)
        norm.weight = pattern(width,56)
        x = pattern(rows*width,55).reshape(1,rows,width)
        emit(f"groupnorm_r{rows}_w{width}_g{group}",[norm(x)],rows=rows,width=width,group=group)
    for tokens in (1,4,7,11,2048):
        for mode in ("none","contiguous","projected"):
            norm = SigmoidRMSNormGated(128,eps=1e-6)
            norm.weight = pattern(128,56)
            x = pattern(tokens*6144,55).reshape(1,tokens,48,128)
            gate = None
            if mode == "contiguous":
                gate = pattern(tokens*6144,57).reshape(1,tokens,48,128)
            elif mode == "projected":
                gate = pattern(tokens*16480,57).reshape(1,tokens,16480)[...,10240:16384].reshape(1,tokens,48,128)
            emit(f"gatednorm_t{tokens}_{mode}",[norm(x,gate)],tokens=tokens,mode=mode)
    for rows in (1,4,7,32,2048):
        for stride in (48,16480):
            for initial in (False,True):
                for masked in (False,True):
                    state = pattern(48*128*128,15).astype(mx.float32).reshape(1,48,128,128) if initial else None
                    a_log,dt_bias = pattern(48,62),pattern(48,63)
                    mask = mx.array((np.arange(rows)%3)!=0).reshape(1,rows) if masked else None
                    for step in range(2):
                        q = (pattern(rows*2048,58+step)/16).reshape(1,rows,16,128)
                        k = (pattern(rows*2048,59+step)/16).reshape(1,rows,16,128)
                        v = pattern(rows*6144,60+step).reshape(1,rows,48,128)
                        if stride==48:
                            a,b = pattern(rows*48,60+step).reshape(1,rows,48),pattern(rows*48,61+step).reshape(1,rows,48)
                        else:
                            projected = pattern(rows*16480,60+step).reshape(1,rows,16480)
                            a,b = projected[...,16432:],projected[...,16384:16432]
                        y,state = gated_delta.gated_delta_update(q,k,v,a,b,a_log,dt_bias,state,mask)
                        emit(f"delta_update_r{rows}_stride{stride}_initial{int(initial)}_mask{int(masked)}_step{step}",
                             [y,state,gated_delta.compute_g(a_log,a,dt_bias),mx.sigmoid(b)],
                             rows=rows,stride=stride,initial=initial,masked=masked,step=step)


if __name__ == '__main__':
    main()