fix(speculative): guard PEAGLE flex attention compile (#2443)

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Author: akoumpa

nemo_automodel/components/speculative/eagle/peagle_draft.py

@@ -36,19 +36,25 @@
 from nemo_automodel.components.speculative.eagle.peagle_attention import create_peagle_mask_mod
 
 # flex_attention compiled for the CUDA training path. Inductor's flex backend is
-# not available on CPU (it raises ``InductorError``), so ``_peagle_flex_attention``
-# dispatches to eager ``flex_attention`` there -- correct, just slower -- which
-# keeps the P-EAGLE unit tests and CPU smoke checks runnable. The compiled
+# not available on CPU, and it currently requires query/key/value head dimensions
+# of at least 16. ``_peagle_flex_attention`` dispatches to eager
+# ``flex_attention`` for unsupported cases -- correct, just slower -- which keeps
+# P-EAGLE unit tests and small CPU/GPU smoke checks runnable. The compiled
 # callable is lazy, so importing this module on CPU costs nothing.
 _peagle_flex_attention_compiled = torch.compile(
     flex_attention,
     mode="max-autotune-no-cudagraphs",
 )
 
 
+def _peagle_compile_supported(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> bool:
+    """Return whether Inductor's flex-attention lowering supports these tensors."""
+    return q.is_cuda and q.shape[-1] >= 16 and k.shape[-1] >= 16 and v.shape[-1] >= 16
+
+
 def _peagle_flex_attention(q, k, v, *, block_mask, scale):
-    """Run the P-EAGLE flex attention, compiling only on CUDA."""
-    flex = _peagle_flex_attention_compiled if q.is_cuda else flex_attention
+    """Run the P-EAGLE flex attention, compiling only when Inductor supports it."""
+    flex = _peagle_flex_attention_compiled if _peagle_compile_supported(q, k, v) else flex_attention
     return flex(q, k, v, block_mask=block_mask, scale=scale)
 
 

tests/unit_tests/speculative/test_peagle.py

@@ -51,16 +51,6 @@
     reason="P-EAGLE forward uses flex_attention, which has no CPU autograd support in the CI torch build.",
 )
 
-# P-EAGLE's draft forward runs flex_attention, whose autograd is not implemented
-# on CPU in the CI torch build (even the forward errors once the inputs require
-# grad). Tests that drive the draft forward therefore run on CUDA and are skipped
-# when it is unavailable; the pure-tensor / mask-mod / config tests stay on CPU.
-_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-_gpu_only = pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="P-EAGLE forward uses flex_attention, which has no CPU autograd support in the CI torch build.",
-)
-
 
 def _build_tiny_draft_model(
     *, parallel_drafting: bool = False, mask_token_id: int = 0, device: str = _DEVICE