[pkmandke]

tl;dr: module_input_spec={0: ...} on a decomposed SDPA module resolved to a SymInt from sym_size. prepare() then attached a FakeQuantize to a Python int and crashed on tensor.ndim. This fix skips non-float boundary values during prepare() in graph mode.

---

A module-level module_input_spec / module_output_spec that resolves to a
shape integer crashes prepare:

AttributeError: 'int' object has no attribute 'ndim'
  qparams_calculator.py:340  self._resolve_axis(tensor.ndim)

The prepared graph has activation_post_process_N(sym_size_int_M): a
FakeQuantize attached to a Python int.

Change/Fix

Skip non-floating-point values in _find_and_apply_module_level_spec (inputs and
outputs).

How it was found

module_input_spec={0: <fp8>} on a decomposed SDPA module, meant to quantize the
query. SDPA builds its causal mask from shape integers before touching the query,
so boundary input 0 ends up being a SymInt.

Tests

TestNonFloatModuleBoundary in test_graph_mode_quantizer.py, for the wildcard
and {0: spec} cases. Both fail before this change with the error above and
pass after.

[vineet-g]

You can fold into TestSingleElementTensors next to test_single_element_module_boundary_not_observed. It uses the same model and config pattern. Broaden that class's docstring to "non-float or ≤1 element".

[pkmandke]

done

[pkmandke]

I think we should also remove "matching eager mode" from the docstring. Such comments can get stale (if eager mode later changes behavior) and can be misleading..

[vineet-g]

assert fq_nodes always passes because the global config quantizes linear. Assert that the float boundary input y is observed instead.

[pkmandke]

Ah good point!

I checked what gets observed and found that asserting y is observed would have the same problem since y is the linear output, which the global activation spec also observes in every case, so it passes with or without the module spec.

I changed the test so the assertion is meaningful instead:

• The config no longer has a global part. It only sets module_input_spec on block, so y can only be observed through the module spec.
• The test asserts that linear (y) is observed.
• The indexed case is now {0: spec, 1: spec} instead of {0: spec}. With only index 0, y is boundary input 1 and is never selected, so the assert would fail.
• The check that every observed node is a float tensor stays. It is what catches the original crash..