trunk/c15e9774278597951aa402693c1bbcb6c8c7b9e8: Fix SAC FIFO mismatch for inductor_compiled_code HOP (#177198)

"""Minimal self-contained repro for DTensor int tensor + SAC + compile bug."""



import os

os.environ["RANK"] = "0"

os.environ["WORLD_SIZE"] = "1"

os.environ["LOCAL_RANK"] = "0"

os.environ["MASTER_ADDR"] = "localhost"

os.environ["MASTER_PORT"] = "29500"



import torch

import torch.distributed as dist

import torch.utils.checkpoint

from functools import partial

from torch.distributed.device_mesh import init_device_mesh

from torch.distributed.tensor import DTensor, Replicate

from torch.nn.attention.flex_attention import create_block_mask, flex_attention

from torch.utils.checkpoint import CheckpointPolicy



torch._inductor.config.wrap_inductor_compiled_regions = True

dist.init_process_group(backend="fake", init_method="env://", world_size=1, rank=0)

device = "cuda" if torch.cuda.is_available() else "cpu"

mesh = init_device_mesh(device, mesh_shape=(1,), mesh_dim_names=("fsdp",))



_SAVE = {torch.ops.higher_order.flex_attention.name(): None, "inductor_compiled_code": None}



def ac_policy(ctx, op, *args, **kwargs):

    if (isinstance(op, torch._ops.HigherOrderOperator) and op.name() in _SAVE) or op in _SAVE:

        return CheckpointPolicy.MUST_SAVE

    return CheckpointPolicy.PREFER_RECOMPUTE



L = 64

_compiled_create_block_mask = torch.compile(create_block_mask, dynamic=False, fullgraph=True)

_compiled_flex_attention = torch.compile(flex_attention, mode="default", fullgraph=True, dynamic=False)



def causal(b, h, q, k):

    return q >= k



_mask_cache = {}



def inner_fn(q, mask):

    ql = q.to_local()

    out = _compiled_flex_attention(query=ql, key=ql, value=ql, block_mask=mask)

    return DTensor.from_local(out, device_mesh=q.device_mesh, placements=q.placements)



inner_fn = torch.compile(inner_fn, mode="default", fullgraph=True, dynamic=False)



def outer_fn(x):

    if "m" not in _mask_cache:

        _mask_cache["m"] = _compiled_create_block_mask(

            mask_mod=causal, B=None, H=None,

            Q_LEN=L, KV_LEN=L, BLOCK_SIZE=128, device=x.device,

        )

    return inner_fn(x, _mask_cache["m"])



context_fn = partial(torch.utils.checkpoint.create_selective_checkpoint_contexts, ac_policy)



x = DTensor.from_local(

    torch.randn(1, 1, L, L, device=device, dtype=torch.bfloat16, requires_grad=True),

    mesh, (Replicate(),),

)



try:

    out = torch.utils.checkpoint.checkpoint(

        outer_fn, x, use_reentrant=False, context_fn=context_fn,

    )

    out.sum().backward()

    print("PASS - no error")

except RuntimeError as e:

    import traceback

    traceback.print_exc()

    if "Only Tensors of floating point and complex dtype" in str(e):

        print("\nFAIL - DTensor constructor error reproduced!")

    else:

        print(f"\nDifferent error: {e}")

finally:

    dist.destroy_process_group()