Fix for the low-probability garbled output issue in the GLM-5 series models. (#22811)
This commit is contained in:
@@ -1887,6 +1887,8 @@ class NSATokenToKVPool(MLATokenToKVPool):
|
||||
seq_len: int,
|
||||
page_indices: torch.Tensor,
|
||||
):
|
||||
if self.layer_transfer_counter is not None:
|
||||
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
|
||||
buf = self.index_k_with_scale_buffer[layer_id - self.start_layer]
|
||||
return index_buf_accessor.GetK.execute(
|
||||
self, buf, seq_len=seq_len, page_indices=page_indices
|
||||
@@ -1898,6 +1900,8 @@ class NSATokenToKVPool(MLATokenToKVPool):
|
||||
seq_len: int,
|
||||
page_indices: torch.Tensor,
|
||||
):
|
||||
if self.layer_transfer_counter is not None:
|
||||
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
|
||||
buf = self.index_k_with_scale_buffer[layer_id - self.start_layer]
|
||||
return index_buf_accessor.GetS.execute(
|
||||
self, buf, seq_len=seq_len, page_indices=page_indices
|
||||
@@ -1922,6 +1926,8 @@ class NSATokenToKVPool(MLATokenToKVPool):
|
||||
k_fp8: (seq_len, index_head_dim), uint8
|
||||
k_scale: (seq_len, 4), uint8
|
||||
"""
|
||||
if self.layer_transfer_counter is not None:
|
||||
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
|
||||
buf = self.index_k_with_scale_buffer[layer_id - self.start_layer]
|
||||
return index_buf_accessor.GetKAndS.execute(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user