From 9a7867539c9f22790677933d5b81e764586f666b Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Wed, 19 Aug 2026 14:15:23 -0700 Subject: [PATCH] test(dsa): compare only the effective top-k slice When top_k exceeds a row's sequence length, the reference pads the rest of the row with zeros while the kernel leaves the slots at their sentinel init, so comparing whole rows compared two different paddings and failed with inf differences. Cap both sides at min(top_k, seq_len), which is what the surrounding comment already described. --- test/python/fe_api/dsa/dsa_reference.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/test/python/fe_api/dsa/dsa_reference.py b/test/python/fe_api/dsa/dsa_reference.py index c5596df7a..863c73a9a 100644 --- a/test/python/fe_api/dsa/dsa_reference.py +++ b/test/python/fe_api/dsa/dsa_reference.py @@ -282,17 +282,19 @@ def check_ref_indexer_top_k( # effective length. The DSA kernel returns indices for a particular row; we # verify that the set of picked indices matches the reference set. n_rows = input_values.shape[0] + batch = seq_lens.shape[0] + rows_per_batch = n_rows // batch for r in range(n_rows): - # Build sets of picked indices, but cap at the smaller of the effective - # lengths. The reference pads with zeros beyond seq_lens[b]; so do - # actual. We rely on value equality at matched indices via the value - # tensor check when return_val is True. - ref_set = set(int(i) for i in idx_ref[r].tolist()) - act_set = set(int(i) for i in idx_actual[r].tolist()) + # Only the first min(top_k, seq_lens[b]) entries of a row are picked. + # Past that the reference pads with zeros while the kernel leaves its + # sentinel init, so comparing whole rows compares two paddings. + k_eff = min(top_k, int(seq_lens[r // rows_per_batch].item())) + ref_set = set(int(i) for i in idx_ref[r, :k_eff].tolist()) + act_set = set(int(i) for i in idx_actual[r, :k_eff].tolist()) # Values within the effective top-k slice should match after sorting. if return_val: - ref_sorted = torch.sort(val_ref[r]).values - act_sorted = torch.sort(val_actual[r].to(val_ref.dtype)).values + ref_sorted = torch.sort(val_ref[r, :k_eff]).values + act_sorted = torch.sort(val_actual[r, :k_eff].to(val_ref.dtype)).values torch.testing.assert_close(act_sorted, ref_sorted, atol=atol, rtol=rtol)