[Test] Split test_sampling_mask_preserves_top_k_boundary_ties to remove Triton-kernel-specific assumption Description (#55315)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
Qiming Zhang committed
Sep 4, 2026 at 06:16 UTC
07ee97e0230e0ff8759cdc95a1bb09ba2e462720
1 file changed
+25
-2
tests/v1/test_outputs.py
+25
-2
@@ -104,7 +104,10 @@ def test_sampling_mask_tensors_match_finite_support(max_num_kept):
104
assert tensors.to_cpu_nonblocking().tolists().to_nested_list() == expected
105
106
107
-def test_sampling_mask_preserves_top_k_boundary_ties():
107
+def test_sampling_mask_matches_processed_top_k_top_p_support():
108
+ """The mask must exactly mirror whatever support `apply_top_k_top_p`
109
+ (the real logits-processing function used by the sampler) actually
110
+ produces, whatever backend implements it."""
111
processed_logits = apply_top_k_top_p(
112
logits=torch.tensor(
113
[[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device=DEVICE_TYPE
@@ -115,7 +118,27 @@ def test_sampling_mask_preserves_top_k_boundary_ties():
118
expected_token_ids = (
119
torch.isfinite(processed_logits[0]).nonzero().flatten().tolist()
120
)
118
- assert len(expected_token_ids) > 3
121
+ assert 0 < len(expected_token_ids) <= processed_logits.shape[1]
122
+
123
+ tensors = SamplingMaskTensors.from_logits(
124
+ processed_logits,
125
+ num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE),
126
+ max_num_kept=3,
127
+ )
128
+ result = tensors.tolists()
129
+
130
+ assert result.to_nested_list() == [expected_token_ids]
131
+
132
+
133
+def test_sampling_mask_preserves_top_k_boundary_ties():
134
+ """When the kept support is wider than `max_num_kept` (e.g. a top-k
135
+ boundary tie keeps more than k logits), the mask must fall back to the
136
+ exact bitmask instead of silently truncating to `max_num_kept` ids."""
137
+ processed_logits = torch.tensor(
138
+ [[6.0, 5.0, 4.0, 4.0, 4.0, float("-inf"), float("-inf"), float("-inf")]],
139
+ device=DEVICE_TYPE,
140
+ )
141
+ expected_token_ids = [0, 1, 2, 3, 4]
142
143
tensors = SamplingMaskTensors.from_logits(
144
processed_logits,