[opt] improved performance

leideng · leideng · commit ae6a60180aa8 · 2026-01-22T16:29:08.000+08:00
diff --git a/ucm/sparse/gsa_on_device/gsa_on_device.py b/ucm/sparse/gsa_on_device/gsa_on_device.py
@@ -382,7 +382,7 @@ def attention_begin(
                                 : self.batch_size_for_hamming
                             ]
                             self.max_seq_len_for_hamming = torch.max(
-                                self.seq_lens_for_hamming
+                                attn_metadata.seq_lens[: self.batch_size_for_hamming]
                             ).item()
                             self.block_table_decode = self.ori_block_table_decode[
                                 : self.batch_size_for_hamming

Original file line number	Diff line number	Diff line change
`@@ -382,7 +382,7 @@ def attention_begin(`
`382`	`382`	`: self.batch_size_for_hamming`
`383`	`383`	`]`
`384`	`384`	`self.max_seq_len_for_hamming = torch.max(`
`385`		`- self.seq_lens_for_hamming`
	`385`	`+ attn_metadata.seq_lens[: self.batch_size_for_hamming]`
`386`	`386`	`).item()`
`387`	`387`	`self.block_table_decode = self.ori_block_table_decode[`
`388`	`388`	`: self.batch_size_for_hamming`