[None][test] Declare supported attention test phases

tests/unittest/_torch/attention/model_attn_config.py

@@ -52,7 +52,9 @@
 """
 
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import Literal, Optional
+
+AttentionPhase = Literal["ctx", "gen"]
 
 
 @dataclass(frozen=True)
@@ -68,11 +70,11 @@ class ModelAttnConfig:
     no_cache: bool = False  # bidirectional DiT/encoder (no KV cache)
     is_cross: bool = False  # encoder-decoder cross attention
     # MLA (DeepSeek-style latent attention). For absorbed generation num_kv_heads
-    # is 1 (single latent head); for the up-projected context pass (mla_context)
-    # it is MHA (num_kv_heads == num_heads) with asymmetric K/V (K head_dim =
-    # qk_nope + qk_rope, V head_dim = v_head_dim).
+    # is 1 (single latent head); for the up-projected context pass it is MHA
+    # (num_kv_heads == num_heads) with asymmetric K/V (K head_dim = qk_nope +
+    # qk_rope, V head_dim = v_head_dim).
     is_mla: bool = False
-    mla_context: bool = False  # up-projected MHA context pass (vs absorbed gen)
+    phases: tuple[AttentionPhase, ...] | None = None
     kv_lora_rank: Optional[int] = None
     q_lora_rank: Optional[int] = None
     qk_nope_head_dim: Optional[int] = None
@@ -576,6 +578,7 @@ class ModelAttnConfig:
         num_heads=128,
         num_kv_heads=1,
         head_dim=192,  # qk_nope+qk_rope
+        phases=("gen",),
         is_mla=True,
         kv_lora_rank=512,
         q_lora_rank=1536,
@@ -590,6 +593,7 @@ class ModelAttnConfig:
         num_heads=32,
         num_kv_heads=1,
         head_dim=192,
+        phases=("gen",),
         is_mla=True,
         kv_lora_rank=512,
         q_lora_rank=1536,
@@ -604,6 +608,7 @@ class ModelAttnConfig:
         num_heads=64,
         num_kv_heads=1,
         head_dim=192,
+        phases=("gen",),
         is_mla=True,
         kv_lora_rank=512,
         q_lora_rank=1536,
@@ -620,7 +625,7 @@ class ModelAttnConfig:
         num_heads=128,
         num_kv_heads=128,
         head_dim=192,
-        mla_context=True,
+        phases=("ctx",),
         is_mla=True,
         kv_lora_rank=512,
         q_lora_rank=1536,
@@ -634,7 +639,7 @@ class ModelAttnConfig:
         num_heads=32,
         num_kv_heads=32,
         head_dim=192,
-        mla_context=True,
+        phases=("ctx",),
         is_mla=True,
         kv_lora_rank=512,
         q_lora_rank=1536,
@@ -648,7 +653,7 @@ class ModelAttnConfig:
         num_heads=64,
         num_kv_heads=64,
         head_dim=192,
-        mla_context=True,
+        phases=("ctx",),
         is_mla=True,
         kv_lora_rank=512,
         q_lora_rank=1536,
@@ -813,4 +818,4 @@ class ModelAttnConfig:
     ),
 ]
 
-MODEL_CONFIGS: List[ModelAttnConfig] = _STANDARD + _MLA + _CROSS + _NO_CACHE
+MODEL_CONFIGS: list[ModelAttnConfig] = _STANDARD + _MLA + _CROSS + _NO_CACHE

tests/unittest/_torch/attention/test_attention_backends.py

@@ -9,7 +9,7 @@
 configurations actually used by the supported models (``model_attn_config.py``), so
 each case maps to a real workload. The orthogonal dimensions are bounded:
 
-* default cross on every cacheable config: phase {ctx, dec, mix} x
+* default cross on every cacheable config: phase {ctx, gen, mix} x
   precision {bf16, fp8-KV} x KV-manager {v1, v2}, at page_size=32, layout=HND.
 * the non-default dimension values (page_size=64, layout=NHD, dtype=fp16) are
   exercised on a small representative set (GQA / MHA / MQA) to avoid a full
@@ -85,6 +85,10 @@ def _phases_for(cfg: ModelAttnConfig) -> dict:
     return _phases_from_window(cfg.sliding_window)
 
 
+def _phases_to_run(cfg: ModelAttnConfig, available_phases: dict) -> tuple[str, ...]:
+    return tuple(available_phases) if cfg.phases is None else cfg.phases
+
+
 def _rope_dict(cfg: ModelAttnConfig):
     if cfg.rope is None:
         return None
@@ -193,13 +197,12 @@ def _expand(cfg: ModelAttnConfig, precisions, kv_layouts, page_sizes):
                                 **base,
                             ),
                         )
-                    elif cfg.is_mla and cfg.mla_context:
-                        yield f"{cfg.id}-ctx-{tag}", BackendCase(**phases["ctx"], **base)
-                    elif cfg.is_mla:
-                        yield f"{cfg.id}-gen-{tag}", BackendCase(**phases["gen"], **base)
                     else:
-                        for phase_name, phase in phases.items():
-                            yield f"{cfg.id}-{phase_name}-{tag}", BackendCase(**phase, **base)
+                        for phase_name in _phases_to_run(cfg, phases):
+                            yield (
+                                f"{cfg.id}-{phase_name}-{tag}",
+                                BackendCase(**phases[phase_name], **base),
+                            )
 
 
 def _model_cases():