From f1abeba881dfb87593fc4de466be7e2aaf61b66c Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Wed, 22 Jul 2026 01:00:56 -0700 Subject: [PATCH] [None][test] Declare supported attention test phases Replace mla_context with phases so each model attention config declares whether it supports context, generation, or both. Keep the existing context, generation, and mixed cases for standard attention when phases is None. Explicit MLA phases select only their declared atomic paths, so a config supporting context and generation does not imply an MLA mixed batch. Use immutable typed tuples for phase declarations. Signed-off-by: Yihan Wang --- .../_torch/attention/model_attn_config.py | 23 +++++++++++-------- .../attention/test_attention_backends.py | 17 ++++++++------ 2 files changed, 24 insertions(+), 16 deletions(-) diff --git a/tests/unittest/_torch/attention/model_attn_config.py b/tests/unittest/_torch/attention/model_attn_config.py index 29f292a89c12..b6f32e2fbbe3 100644 --- a/tests/unittest/_torch/attention/model_attn_config.py +++ b/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 diff --git a/tests/unittest/_torch/attention/test_attention_backends.py b/tests/unittest/_torch/attention/test_attention_backends.py index 8bae3401ba7a..ca00dcd21125 100644 --- a/tests/unittest/_torch/attention/test_attention_backends.py +++ b/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():