b12x CompilePool: compiling in a worker that already described the job publishes nothing (lru_cache in mla_compress escapes evict_planning_artifacts)

Start with b12x/attention/mla_compress/_cute.py and b12x/_lib/program_cache.py to trace how compile_compress is memoized and which caches evict_planning_artifacts() clears. Then inspect the _run_job path in b12x/_lib/compile_pool.py and the other memoized compile helpers mentioned in the issue. Reproduce with the supplied one-worker script; done means the job publishes its expected artifact after describing and compiling in the same worker.

Summary

In b12x's CompilePool, a compile job that runs in a worker which has already described the same job returns successfully but publishes nothing to B12X_COMPILE_CACHE_DIR. The cause is compile_compress in b12x/attention/mla_compress/_cute.py. It is memoized with a plain functools.lru_cache, so describing the job stores the planning placeholder in that memo. evict_planning_artifacts() only clears caches registered through b12x._lib.program_cache, so it never removes that entry. The later _run_job in the same process gets the placeholder back, the program keys still match, and nothing is compiled.

The same pattern exists in flashinfer-ai/flashinfer main (flashinfer/experimental/b12x/attention/mla_compress/_cute.py:277). Other compile helpers also memoize with lru_cache, for example norm/vision.py:125 and comm/pcie/_dma_kernels.py. We have not checked whether they behave the same way.

Reproduction

Environment: Flare ca1f960e with its pinned FlashInfer 344e4837, torch 2.14.1+cu132, NVIDIA GB10. The script uses only CompilePool with one worker, so compile() describes and compiles in the same process. The worker never initializes CUDA.

"""CompilePool.compile() reports success but publishes no artifact for mla_compress.

A one-worker pool describes the job (CompilePool.plan -> _describe_jobs) and then compiles
it (_run_job) in the same process. compile_compress is memoized with functools.lru_cache,
which evict_planning_artifacts() does not reach, so the compile returns the planning
placeholder and nothing is written to B12X_COMPILE_CACHE_DIR.
"""
import os
import sys
import tempfile

if "B12X_COMPILE_CACHE_DIR" not in os.environ:  # spawned workers inherit the parent's choice
    root = tempfile.mkdtemp(prefix="b12x-repro-")
    os.environ.update(B12X_COMPILE_CACHE_DIR=f"{root}/b12x", TRITON_CACHE_DIR=f"{root}/triton")

import torch
from b12x._lib.compile_pool import CompileJob, CompilePool
from b12x._lib.compile_plan import compiled_program_available

if os.environ.get("B12X_REPRO_FIX") == "1":
    # Candidate fix: memoize with b12x's program_cache, which evict_planning_artifacts() reaches.
    # Module scope runs in spawned workers too, so they compile with the patched factory.
    from b12x._lib.program_cache import program_cache
    from b12x.attention.mla_compress import _cute

    _cute.compile_compress = program_cache(_cute.compile_compress.__wrapped__)

JOB = CompileJob.create(
    "b12x.attention.mla_compress._preparation:compile_mla_compress",
    {"max_tokens": 2560, "max_states": 1, "max_requests": 1, "ratio": 1, "head_dim": 512},
    CompileJob.DEVICE_ORDINAL,
)


def run(workers, describe_in_worker=True):
    properties = torch.cuda.get_device_properties(0)
    with CompilePool(
        device_ordinal=0,
        compute_capability=(properties.major, properties.minor),
        device_uuid=str(properties.uuid),
        product_name=properties.name,
        sm_count=properties.multi_processor_count,
        max_shared_memory_per_block=properties.shared_memory_per_block_optin,
        max_shared_memory_per_multiprocessor=properties.shared_memory_per_multiprocessor,
        workers=workers,
    ) as pool:
        if describe_in_worker:
            summary = pool.compile([JOB])  # plan() describes in the worker, then it compiles there
        else:
            from b12x._lib.compile_pool import describe_compilation
            plan = describe_compilation(JOB)  # describe in this process; the worker only compiles
            pool.wait(pool.submit_plans([plan]))
            summary = pool.summary()
        programs = sorted(pool.programs, key=lambda p: (p.dialect, p.key))
    return summary, [(p.name, compiled_program_available(p)) for p in programs]


if __name__ == "__main__":
    print(f"torch {torch.__version__}, {torch.cuda.get_device_name(0)}")
    mode = sys.argv[1] if len(sys.argv) > 1 else "worker"
    summary, published = run(workers=1, describe_in_worker=mode == "worker")
    print(f"describe in {mode}: jobs={summary.jobs} cute_compilations={summary.cute_compilations} "
          f"published={published}")
    sys.exit(0 if all(ok for _, ok in published) else 1)

$ python repro.py worker        # CompilePool.compile(): describe and compile in the same worker
describe in worker: jobs=1 cute_compilations=0 published=[('attention.mla_compress.cute', False)]
$ python repro.py parent        # describe in the parent, the worker only compiles
describe in parent: jobs=1 cute_compilations=1 published=[('attention.mla_compress.cute', True)]
$ B12X_"""CompilePool.compile() reports success but publishes no artifact for mla_compress.

A one-worker pool describes the job (CompilePool.plan -> _describe_jobs) and then compiles
it (_run_job) in the same process. compile_compress is memoized with functools.lru_cache,
which evict_planning_artifacts() does not reach, so the compile returns the planning
placeholder and nothing is written to B12X_COMPILE_CACHE_DIR.
"""
import os
import sys
import tempfile

if "B12X_COMPILE_CACHE_DIR" not in os.environ:  # spawned workers inherit the parent's choice
    root = tempfile.mkdtemp(prefix="b12x-repro-")
    os.environ.update(B12X_COMPILE_CACHE_DIR=f"{root}/b12x", TRITON_CACHE_DIR=f"{root}/triton")

import torch
from b12x._lib.compile_pool import CompileJob, CompilePool
from b12x._lib.compile_plan import compiled_program_available

if os.environ.get("B12X_REPRO_FIX") == "1":
    # Candidate fix: memoize with b12x's program_cache, which evict_planning_artifacts() reaches.
    # Module scope runs in spawned workers too, so they compile with the patched factory.
    from b12x._lib.program_cache import program_cache
    from b12x.attention.mla_compress import _cute

    _cute.compile_compress = program_cache(_cute.compile_compress.__wrapped__)

JOB = CompileJob.create(
    "b12x.attention.mla_compress._preparation:compile_mla_compress",
    {"max_tokens": 2560, "max_states": 1, "max_requests": 1, "ratio": 1, "head_dim": 512},
    CompileJob.DEVICE_ORDINAL,
)


def run(workers, describe_in_worker=True):
    properties = torch.cuda.get_device_properties(0)
    with CompilePool(
        device_ordinal=0,
        compute_capability=(properties.major, properties.minor),
        device_uuid=str(properties.uuid),
        product_name=properties.name,
        sm_count=properties.multi_processor_count,
        max_shared_memory_per_block=properties.shared_memory_per_block_optin,
        max_shared_memory_per_multiprocessor=properties.shared_memory_per_multiprocessor,
        workers=workers,
    ) as pool:
        if describe_in_worker:
            summary = pool.compile([JOB])  # plan() describes in the worker, then it compiles there
        else:
            from b12x._lib.compile_pool import describe_compilation
            plan = describe_compilation(JOB)  # describe in this process; the worker only compiles
            pool.wait(pool.submit_plans([plan]))
            summary = pool.summary()
        programs = sorted(pool.programs, key=lambda p: (p.dialect, p.key))
    return summary, [(p.name, compiled_program_available(p)) for p in programs]


if __name__ == "__main__":
    print(f"torch {torch.__version__}, {torch.cuda.get_device_name(0)}")
    mode = sys.argv[1] if len(sys.argv) > 1 else "worker"
    summary, published = run(workers=1, describe_in_worker=mode == "worker")
    print(f"describe in {mode}: jobs={summary.jobs} cute_compilations={summary.cute_compilations} "
          f"published={published}")
    sys.exit(0 if all(ok for _, ok in published) else 1)_FIX=1 python repro.py worker   # compile_compress memoized with program_cache instead
describe in worker: jobs=1 cute_compilations=1 published=[('attention.mla_compress.cute', True)]

In the first case, compile() returns with no error and the summary counts the job, but zero CuTe compilations ran and the program is missing from the cache.

Where we hit it

We were testing kernel precompilation for DeepSeek V4.1 Flash TP3 on three DGX Sparks. The candidate compile jobs were recorded, then described and compiled in a separate many-worker pool before the export loads the model. With 20 workers per host, 16 of the 96 distinct attention.mla_compress.cute programs were never published. Each of those programs has three equivalent payloads, so about 3 of the 20 workers held a stale memo entry for it. Compiling in fresh workers after describing, or applying the fix below, removes the gap.

Stock preparation also describes candidates through the pool (_discover_candidates, preparation/session.py:1658) and then compiles in the same workers. In our stock exports we did not see a failure from this, so we cannot say whether any current family reaches it there. If it does, _wait_programs should report "compiler completed without publishing required artifacts" rather than leave a silent gap.

Suggested fix

Memoize compile_compress with b12x._lib.program_cache.program_cache instead of functools.lru_cache, so that evict_planning_artifacts() reaches it. That is the third case above. The other lru_cache compile helpers may need the same change.

A second, more defensive option: _run_job could check compiled_program_available() for its expected programs after job.invoke() and raise if any is missing, so a stale memo shows up as an error rather than a silent no-op.