[Bugfix] BailingMoeV3 KDA: skip absent metadata during CUDA graph profiling (#53593)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-08-25 01:02:50 +00:00
parent e3dde1ee9b
commit 09fddeb4ce
1 changed files with 4 additions and 1 deletions
@@ -813,7 +813,10 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase):
            return

        assert isinstance(attn_metadata_map, dict)
-        attn_metadata = attn_metadata_map[self.prefix]
+        attn_metadata = attn_metadata_map.get(self.prefix)
+        if attn_metadata is None:
+            # Profile/warmup dummy runs skip mamba-family metadata.
+            return
        assert isinstance(attn_metadata, GDNAttentionMetadata)
        has_initial_state = attn_metadata.has_initial_state
        spec_query_start_loc = attn_metadata.spec_query_start_loc