Fix for the low-probability garbled output issue in the GLM-5 series models. (#22811)

This commit is contained in:
Yuxuan Zhang
2026-04-17 09:52:13 +08:00
committed by GitHub
parent 0052093178
commit 16d11c2a10
@@ -1887,6 +1887,8 @@ class NSATokenToKVPool(MLATokenToKVPool):
seq_len: int, seq_len: int,
page_indices: torch.Tensor, 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] buf = self.index_k_with_scale_buffer[layer_id - self.start_layer]
return index_buf_accessor.GetK.execute( return index_buf_accessor.GetK.execute(
self, buf, seq_len=seq_len, page_indices=page_indices self, buf, seq_len=seq_len, page_indices=page_indices
@@ -1898,6 +1900,8 @@ class NSATokenToKVPool(MLATokenToKVPool):
seq_len: int, seq_len: int,
page_indices: torch.Tensor, 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] buf = self.index_k_with_scale_buffer[layer_id - self.start_layer]
return index_buf_accessor.GetS.execute( return index_buf_accessor.GetS.execute(
self, buf, seq_len=seq_len, page_indices=page_indices 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_fp8: (seq_len, index_head_dim), uint8
k_scale: (seq_len, 4), 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] buf = self.index_k_with_scale_buffer[layer_id - self.start_layer]
return index_buf_accessor.GetKAndS.execute( return index_buf_accessor.GetKAndS.execute(
self, self,