[XPU] Fix NemotronH (hybrid mamba2) launch on --device xpu (#32227)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
Juan Muneton
2026-08-12 13:23:23 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent 256981ce16
commit cac3269305
3 changed files with 15 additions and 12 deletions
@@ -2303,6 +2303,13 @@ class MHATokenToKVPool(KVCache):
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
return self._get_value_buffer(layer_id)
def get_v_head_dim(self):
# Every layer in this pool is full-attention, so the value head dim is
# uniform and known at construction. Mirrors HybridLinearKVPool's
# get_v_head_dim() so the TritonAttnBackend mambaish branch works when a
# mamba2 config is served by a plain MHA pool (no per-linear-layer split).
return self.v_head_dim
def get_kv_buffer(self, layer_id: int):
return self.get_key_buffer(layer_id), self.get_value_buffer(layer_id)