AutoParallel(...).optimize_placement() crashes when the traced model contains aten._scaled_dot_product_cudnn_attention.default.
This happens with any HuggingFace transformers causal-LM that uses attn_implementation="sdpa" (default) on cu130 / Blackwell. The
assertion is raised inside PyTorch's DTensor strategy code, but it blocks autoparallel for all SDPA-based HF models.
Repro:
import os, torch, torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.placement_types import Shard
from transformers import AutoConfig, AutoModelForCausalLM
from autoparallel.api import AutoParallel
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
dist.init_process_group("nccl")
mesh = init_device_mesh("cuda", (dist.get_world_size(),), mesh_dim_names=("tp",))
config = AutoConfig.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
config.use_cache = False
with torch.device("meta"):
model = AutoModelForCausalLM.from_config(config, attn_implementation="sdpa")
model.init_weights = lambda: None
def input_fn():
return torch.randint(0, config.vocab_size, (2, 32), device="cuda")
with AutoParallel(model, input_fn, mesh) as autop:
autop.add_input_constraints([(Shard(0),)])
autop.add_output_constraints([(Shard(0),)])
autop.optimize_placement() # ← crashes here
Same crash with Qwen/Qwen2.5-0.5B-Instruct.
Stack trace
File ".../autoparallel/optimize_sharding.py", line 326, in build_sharding_metadata
strats[node] = get_placement_options_for_node(...)
File ".../autoparallel/shardings/placement_options.py", line 327, in get_placement_options
out_strat = get_op_strategy(op, op_schema)
File ".../autoparallel/shardings/dtensor_sharding_helpers.py", line 359, in get_op_strategy
return propagator.op_strategy_funcs[op](op_schema)
File ".../torch/distributed/tensor/_ops/_matrix_ops.py", line 1012,
in scaled_dot_product_cudnn_attention_strategy
return expand_to_full_mesh_op_strategy(
File ".../torch/distributed/tensor/_ops/utils.py", line 509,
in expand_to_full_mesh_op_strategy
AssertionError: input_specs(3) != strategies(4: 4 args + 0 kwargs)
Root cause (upstream — torch/distributed/tensor/_ops/_matrix_ops.py, _scaled_dot_product_cudnn_attention_base_strategies):
The strategy table has three rows: all_replicate, num_heads_dim_sharding, batch_dim_sharding. Only the first one conditionally
appends a placement for the optional attn_bias input:
if has_attn_bias:
all_replicate.append(Replicate()) # attn bias ← only this row
The other two rows are hardcoded to 12 entries (9 outputs + q/k/v) and never grow. When the model actually passes attn_bias,
input_args_strategy has length 4 but those rows only declare 3 input placements → the assertion in
expand_to_full_mesh_op_strategy fires.
AutoParallel(...).optimize_placement() crashes when the traced model contains
aten._scaled_dot_product_cudnn_attention.default.This happens with any HuggingFace transformers causal-LM that uses attn_implementation="sdpa" (default) on cu130 / Blackwell. The
assertion is raised inside PyTorch's DTensor strategy code, but it blocks autoparallel for all SDPA-based HF models.
Repro:
Same crash with Qwen/Qwen2.5-0.5B-Instruct.
Stack trace
Root cause (upstream — torch/distributed/tensor/_ops/_matrix_ops.py, _scaled_dot_product_cudnn_attention_base_strategies):
The strategy table has three rows: all_replicate, num_heads_dim_sharding, batch_dim_sharding. Only the first one conditionally
appends a placement for the optional attn_bias input:
The other two rows are hardcoded to 12 entries (9 outputs + q/k/v) and never grow. When the model actually passes attn_bias,
input_args_strategy has length 4 but those rows only declare 3 input placements → the assertion in
expand_to_full_mesh_op_strategy fires.