[Bugfix][ROCm] Change AttentionCGSuppoort in TritonMLA to UNIFORM_SINGLE_TOKEN_DECODE (#47144)
Signed-off-by: Dino Music <Dino.Music@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
music-dino committed
Jul 9, 2026 at 04:09 UTC
2c17d33f4291a55b447317640c81eb61077b1b00
1 file changed
+3
-1
vllm/v1/attention/backends/mla/triton_mla.py
+3
-1
@@ -48,7 +48,9 @@ def _compute_num_kv_splits(max_seq_len: int, sm_count: int) -> int:
48
49
50
class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
51
- _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
51
+ _cudagraph_support: ClassVar[AttentionCGSupport] = (
52
+ AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
53
+ )
54
55
def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
56
super().__init__(kv_cache_spec, layer_names, vllm_config, device)