add stable=True to argsort

vadiklyutiy · vadiklyutiy · commit 80ecbbfef8b9 · 2025-10-08T22:24:38.000+04:00
Signed-off-by: Vadim Gimpelson &lt;vadim.gimpelson@gmail.com&gt;
diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py
@@ -229,7 +229,7 @@ def build(  # type: ignore[override]
                 spec_token_masks = torch.repeat_interleave(
                     spec_sequence_masks, query_lens
                 )
-                indx = torch.argsort(spec_token_masks)
+                indx = torch.argsort(spec_token_masks, stable=True)
                 num_non_spec_tokens = num_prefill_tokens + num_decode_tokens
                 non_spec_token_indx = indx[:num_non_spec_tokens]
                 spec_token_indx = indx[num_non_spec_tokens:]

Original file line number	Diff line number	Diff line change
`@@ -229,7 +229,7 @@ def build( # type: ignore[override]`
`229`	`229`	`spec_token_masks = torch.repeat_interleave(`
`230`	`230`	`spec_sequence_masks, query_lens`
`231`	`231`	`)`
`232`		`- indx = torch.argsort(spec_token_masks)`
	`232`	`+ indx = torch.argsort(spec_token_masks, stable=True)`
`233`	`233`	`num_non_spec_tokens = num_prefill_tokens + num_decode_tokens`
`234`	`234`	`non_spec_token_indx = indx[:num_non_spec_tokens]`
`235`	`235`	`spec_token_indx = indx[num_non_spec_tokens:]`