[XPU] Fix NemotronH (hybrid mamba2) launch on --device xpu (#32227)
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
co-authored by
Ma Mingfei
parent
256981ce16
commit
cac3269305
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user