[ROCm][MLA] Fix mypy: narrow rotary_emb/slot_mapping on the fused MLA path

The new rotary_emb: nn.Module | None on the base MLAAttention broke mypy
(3.10-3.13): guard rotary_emb before use in the fused decode branch, narrow
forward_context.slot_mapping to dict, ignore the impl attr assignment, and
assert rotary_emb in the DeepseekV32 (nvidia) subclass that inherits the type.

co-authored by Claude

Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com>

Author: xaguilar-amd

vllm/model_executor/layers/attention/mla_attention.py

@@ -513,7 +513,7 @@ def __init__(
         self.rotary_emb = rotary_emb
         self._fused_rope_cos_sin: tuple[torch.Tensor, torch.Tensor] | None = None
         if getattr(self.impl, "use_fused_qk_rope_cache", False) and rotary_emb is None:
-            self.impl.use_fused_qk_rope_cache = False
+            self.impl.use_fused_qk_rope_cache = False  # type: ignore[attr-defined]
 
         vllm_config = get_current_vllm_config()
         parallel_config = vllm_config.parallel_config
@@ -900,8 +900,13 @@ def forward_impl(
                         "forward positions when impl.use_fused_qk_rope_cache is set."
                     )
                 forward_context = get_forward_context()
-                slot_mapping = forward_context.slot_mapping[self.layer_name].flatten()
+                fc_slot_mapping = forward_context.slot_mapping
+                assert isinstance(fc_slot_mapping, dict), (
+                    f"Expected slot_mapping to be a dict, got {type(fc_slot_mapping)}."
+                )
+                slot_mapping = fc_slot_mapping[self.layer_name].flatten()
                 cos_cache, sin_cache = self._get_fused_rope_cos_sin()
+                assert self.rotary_emb is not None
                 mqa_q = self.impl.fused_qk_rope_concat_and_cache(  # type: ignore[attr-defined]
                     self,
                     mqa_ql_nope,

vllm/models/deepseek_v32/nvidia/attention.py

@@ -426,6 +426,7 @@ def _fused_attention(
             mla_kv_cache = self.kv_cache
             mla_k_scale = self._k_scale
 
+        assert self.rotary_emb is not None
         q_c = fused_norm_rope(
             positions,
             q_c,