Reject recompute in Eagle TTT, drop wasteful GQA expansion, fix markers

Signed-off-by: seonjinn <sna@nvidia.com>

Author: seonjinn

nemo_rl/models/megatron/draft/eagle_ttt.py

@@ -486,7 +486,12 @@ def from_trunk(
         max_passes: int,
         activation_budget_bytes: int,
     ) -> EagleTTTState:
-        """Validate the complete bound before retaining any supplied tensor."""
+        """Validate the per-layer multi-pass K/V bound before retaining tensors.
+
+        The complete bound across layers, hidden taps, rope, and loss rows is
+        enforced by the provider's storage plan and, at runtime, by the
+        resource ledger; this check covers only this layer's K/V retention.
+        """
         _validate_kv_pair(trunk_key, trunk_value, name="trunk")
         EagleTTTStoragePlan(
             batch_size=trunk_key.shape[0],
@@ -563,6 +568,16 @@ def __init__(
         if attention_type != "self":
             raise ValueError("EAGLE TTT only supports self attention")
         self.layer_number = layer_number
+        recompute_granularity = getattr(config, "recompute_granularity", None)
+        if recompute_granularity is not None:
+            # Activation recompute re-enters this forward during backward, which
+            # violates the one-shot begin_pass/forward/finish_pass lifecycle and
+            # would corrupt the retained multi-pass state on re-entry.
+            raise ValueError(
+                "EAGLE TTT is incompatible with activation recompute "
+                f"(recompute_granularity={recompute_granularity!r}); disable "
+                "recompute for the draft decoder stack"
+            )
         self.context_parallel_size = int(getattr(config, "context_parallel_size", 1))
         self.sequence_parallel = bool(getattr(config, "sequence_parallel", False))
         self.softmax_scale = softmax_scale
@@ -1509,7 +1524,8 @@ def eagle_ttt_attention(
             raise ValueError("sequence layout must match query batch and sequence axes")
         if layout.valid_tokens.device != query.device:
             raise ValueError("sequence layout and query must share a device")
-    _expand_gqa(state.trunk_key, query_heads=query.shape[1])
+    if query.shape[1] % state.trunk_key.shape[1] != 0:
+        raise ValueError("query heads must be divisible by key/value heads")
 
     attention_scale = scale if scale is not None else 1.0 / math.sqrt(query.shape[-1])
     if query.is_cuda and query.dtype != torch.float64:

tests/unit/distributed/test_eagle_ttt_sp.py

@@ -587,6 +587,7 @@ def _sp_packed_params_mismatch_worker(
         dist.destroy_process_group()
 
 
+@pytest.mark.mcore
 @pytest.mark.skipif(
     not torch.distributed.is_nccl_available() or torch.cuda.device_count() < _TP_SIZE,
     reason="two CUDA devices and NCCL are required",
@@ -606,6 +607,7 @@ def test_real_mcore_tp2_sequence_parallel_matches_global_oracle(
     )
 
 
+@pytest.mark.mcore
 @pytest.mark.skipif(
     not torch.distributed.is_nccl_available() or torch.cuda.device_count() < _TP_SIZE,
     reason="two CUDA devices and NCCL are required",
@@ -623,6 +625,7 @@ def test_real_mcore_tp2_sequence_parallel_packed_mismatch_agrees_and_resets(
     )
 
 
+@pytest.mark.mcore
 @pytest.mark.skipif(
     not torch.distributed.is_nccl_available() or torch.cuda.device_count() < _TP_SIZE,
     reason="two CUDA devices and NCCL are required",

tests/unit/models/megatron/test_eagle_ttt_mcore_session.py

@@ -115,6 +115,25 @@ def test_core_preserves_mcore_tp_group_from_process_group_collection() -> None:
     assert core.tp_group is tp_group
 
 
+def test_core_rejects_activation_recompute_config() -> None:
+    # Recompute re-enters the armed forward during backward, which the
+    # one-shot pass lifecycle cannot tolerate; construction must fail fast.
+    module = _load_module()
+    for granularity in ("full", "selective"):
+        with pytest.raises(ValueError, match="recompute"):
+            module.EagleTTTCoreAttention(
+                config=SimpleNamespace(
+                    context_parallel_size=1, recompute_granularity=granularity
+                ),
+                layer_number=1,
+                attn_mask_type=None,
+                attention_type="self",
+                cp_comm_type=None,
+                softmax_scale=None,
+                pg_collection=None,
+            )
+
+
 def test_layer_spec_adapter_is_construction_time_and_does_not_mutate_default() -> None:
     module = _load_module()
     original_core = object()