[AMD][DFlash] Enable Fused KV Materialization (#27854)

Co-authored-by: jonahbernard <Jonah.Bernard@amd.com>
This commit is contained in:
Jonah Bernard
2026-06-12 16:57:46 -07:00
committed by GitHub
co-authored by jonahbernard
parent 3a1417a0c1
commit fda7955890
@@ -238,7 +238,7 @@ class DFlashWorkerV2(BaseSpecWorker):
self._draft_greedy_rank_index_buf: Optional[torch.Tensor] = None
self._draft_greedy_selected_ids_buf: Optional[torch.Tensor] = None
self._draft_greedy_index_cap: int = 0
self._use_fused_kv_materialize = is_cuda()
self._use_fused_kv_materialize = is_cuda() or is_hip()
self._fused_kv_helper: Optional[object] = None
if self._use_fused_kv_materialize:
self._init_fused_kv_helper()