{"record":{"id":"fe2d8cf94e0ccec6","repo":"jax-ml/jax","slug":"async-load-only-supports-tmem-refs","errorCode":null,"errorMessage":"Async load only supports TMEM refs","messagePattern":"Async load only supports TMEM refs","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":4073,"sourceCode":"  )\n  if reduce is None:\n    [result] = results\n    if layout is not None:\n      result = gpu_core.layout_cast(result, layout)\n    return result\n  else:\n    result, reduced = results\n    if layout is not None:\n      result = gpu_core.layout_cast(result, layout)\n    return result, reduced\n\n\n@async_load_tmem_p.def_effectful_abstract_eval\ndef _async_load_tmem_abstract_eval(\n    src, *avals_flat, tree, reduce: tcgen05.LoadReduceOp | None = None\n):\n  if src.memory_space != gpu_core.MemorySpace.TMEM:\n    raise ValueError(\"Async load only supports TMEM refs\")\n  val_aval, effects = state_primitives._get_abstract_eval(\n      src, *avals_flat, tree=tree\n  )\n  if reduce is None:\n    return (val_aval,), effects\n  if val_aval.dtype not in map(jnp.dtype, (jnp.float32, jnp.int32, jnp.uint32)):\n    raise ValueError(\n        f\"Unsupported dtype for reduction: {val_aval.dtype}. Only float32, \"\n        \" int32 and uint32 are supported.\"\n    )\n  reduced_aval = jax_core.ShapedArray(val_aval.shape[:-1], val_aval.dtype)\n  return (val_aval, reduced_aval), effects\n\n\n@lowering.register_lowering_rule(async_load_tmem_p, mgpu.LoweringSemantics.Lane)\ndef _async_load_tmem_lowering_rule(\n    ctx: lowering.LoweringRuleContext,\n    x_ref,","sourceCodeStart":4055,"sourceCodeEnd":4091,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L4055-L4091","documentation":"The async_load_tmem Pallas primitive only accepts references in Tensor Memory (TMEM). Passing a ref whose memory_space is not gpu_core.MemorySpace.TMEM raises ValueError in the abstract eval.","triggerScenarios":"Calling pallas' async TMEM load API (e.g. pl_mosaic.async_load_tmem) on an SMEM or normal (HBM/VMEM) reference.","commonSituations":"Copy-pasting a TMEM kernel snippet onto a regular grid ref; forgetting to allocate the ref in TMEM for Blackwell tcgen05 usage.","solutions":["Allocate the source ref in Tensor Memory (tcgen05 TMEM allocation) before async loading","Use the regular async_load for SMEM and other memory spaces","Check ref.memory_space before calling to fail fast"],"exampleFix":"// before\nval = async_load_tmem(smem_ref)\n// after\nval = async_load_tmem(tmem_ref)  # ref allocated in TMEM","handlingStrategy":"type-guard","validationCode":"from jax._src.pallas import gpu_core\nassert src.memory_space == gpu_core.MemorySpace.TMEM","typeGuard":"def is_tmem(ref): return getattr(ref, 'memory_space', None) == gpu_core.MemorySpace.TMEM","tryCatchPattern":null,"preventionTips":["Allocate tcgen05 TMEM refs explicitly for Blackwell kernels"],"tags":["jax","pallas","tmem","memory-space","blackwell"],"backgroundTag":"invalid-memory-space","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}