Function sample_compiled

src/fused_mm_sampling/core.py:160–166

(*args, seed: int | None = None, **kwargs)

greedy_sample_compiled = nvtx.annotate()(torch.compile(greedy_sample, fullgraph=True))

_sample_compiled = torch.compile(sample, fullgraph=True)


@nvtx.annotate()
def sample_compiled(*args, seed: int | None = None, **kwargs):
    # torch.manual_seed is dynamo-skipped, so seed handling has to live outside
    # the compiled region. Pass seed=None into the compiled inner; the dead
    # `if seed is not None` branch in sample() is folded at trace time.
    if seed is not None:
        torch.manual_seed(seed)