vllm - 💡(How to fix) Fix [CI Failure]: tests/models/language/pooling/test_splade_sparse_pooler.py [3 comments, 2 participants]

[2026-03-30T18:22:59Z] ____________ test_splade_pooler_matches_reference_formula[2-3-5-7] _____________

[2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] B = 2, T = 3, H = 5, V = 7 [2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] @pytest.mark.parametrize("B,T,H,V", [(2, 3, 5, 7)]) [2026-03-30T18:22:59Z] @torch.inference_mode [2026-03-30T18:22:59Z] def test_splade_pooler_matches_reference_formula(B, T, H, V): [2026-03-30T18:22:59Z] """Ensure SPLADESparsePooler forward() matches the mathematical formula: [2026-03-30T18:22:59Z] log1p(relu(logits)) -> max over sequence length (after masking).""" [2026-03-30T18:22:59Z] torch.manual_seed(0) [2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] # Prepare [B] sequences of shape [T, H] [2026-03-30T18:22:59Z] hs_list = [torch.randn(T, H) for _ in range(B)] [2026-03-30T18:22:59Z] hs_tenser = torch.cat(hs_list) [2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] # Simulate PoolingMetadata (only required fields) [2026-03-30T18:22:59Z] prompt_lens = [T, T - 1] [2026-03-30T18:22:59Z] prompt_lens_tenser = torch.tensor(prompt_lens, dtype=torch.int32) [2026-03-30T18:22:59Z] token_ids = torch.tensor( [2026-03-30T18:22:59Z] [ [2026-03-30T18:22:59Z] [101, 5, 102], # Batch 0: [CLS], token, [SEP] [2026-03-30T18:22:59Z] [101, 6, 6], # Batch 1: [CLS], token, token (last token ignored) [2026-03-30T18:22:59Z] ], [2026-03-30T18:22:59Z] dtype=torch.long, [2026-03-30T18:22:59Z] ) [2026-03-30T18:22:59Z] meta = types.SimpleNamespace( [2026-03-30T18:22:59Z] prompt_lens=prompt_lens_tenser, prompt_token_ids=token_ids [2026-03-30T18:22:59Z] ) [2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] # MLM head (prefer BertMLMHead, fallback to Linear if unavailable) [2026-03-30T18:22:59Z] try: [2026-03-30T18:22:59Z] mlm_head = BertMLMHead(hidden_size=H, vocab_size=V, layer_norm_eps=1e-12) [2026-03-30T18:22:59Z] except Exception: [2026-03-30T18:22:59Z] mlm_head = nn.Linear(H, V, bias=True) [2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] # Forward pass through SPLADE pooler [2026-03-30T18:22:59Z] pooler = SPLADESparsePooler(mlm_head=mlm_head, pooling="max", remove_cls_sep=True) [2026-03-30T18:22:59Z] > pooled = pooler(hidden_states=hs_tenser, pooling_metadata=meta) # list of [V] [2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] models/language/pooling/test_splade_sparse_pooler.py:53: [2026-03-30T18:22:59Z] _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ [2026-03-30T18:22:59Z] /usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py:1776: in _wrapped_call_impl [2026-03-30T18:22:59Z] return self._call_impl(*args, **kwargs) [2026-03-30T18:22:59Z] /usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py:1787: in _call_impl [2026-03-30T18:22:59Z] return forward_call(*args, **kwargs) [2026-03-30T18:22:59Z] _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ [2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] self = SPLADESparsePooler( [2026-03-30T18:22:59Z] (mlm_head): BertMLMHead( [2026-03-30T18:22:59Z] (dense): Linear(in_features=5, out_features=5, bias=True) [2026-03-30T18:22:59Z] (acti...yerNorm((5,), eps=1e-12, elementwise_affine=True) [2026-03-30T18:22:59Z] (decoder): Linear(in_features=5, out_features=7, bias=True) [2026-03-30T18:22:59Z] ) [2026-03-30T18:22:59Z] ) [2026-03-30T18:22:59Z] hidden_states = tensor([[ 1.5410, -0.2934, -2.1788, 0.5684, -1.0845], [2026-03-30T18:22:59Z] [-1.3986, 0.4033, 0.8380, -0.7193, -0.4033], [2026-03-30T18:22:59Z] ...1.5091], [2026-03-30T18:22:59Z] [ 0.3704, 1.4565, 0.9398, 0.7748, 0.1919], [2026-03-30T18:22:59Z] [ 1.2638, -1.2904, -0.7911, -0.0209, -0.7185]]) [2026-03-30T18:22:59Z] pooling_metadata = namespace(prompt_lens=tensor([3, 2], dtype=torch.int32), prompt_token_ids=tensor([[101, 5, 102], [2026-03-30T18:22:59Z] [101, 6, 6]])) [2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] def forward( [2026-03-30T18:22:59Z] self, [2026-03-30T18:22:59Z] hidden_states: torch.Tensor, [2026-03-30T18:22:59Z] pooling_metadata: PoolingMetadata, [2026-03-30T18:22:59Z] ) -> SequencePoolerOutput: [2026-03-30T18:22:59Z] lens_tensor = pooling_metadata.prompt_lens [2026-03-30T18:22:59Z] lens: list[int] = lens_tensor.tolist() [2026-03-30T18:22:59Z] B: int = len(lens) [2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] > prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() [2026-03-30T18:22:59Z] E AttributeError: 'types.SimpleNamespace' object has no attribute 'get_prompt_token_ids_cpu' [2026-03-30T18:22:59Z] [2026-03-30T18:22:59Z] /usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/bert.py:641: AttributeError

Fix Plan

The issue arises from the pooling_metadata object not having a get_prompt_token_ids_cpu method. To fix this, we need to ensure that the pooling_metadata object has this method or modify the code to not require it.

Step 1: Modify the SPLADESparsePooler class

We can modify the SPLADESparsePooler class to directly access the prompt_token_ids attribute from the pooling_metadata object.

class SPLADESparsePooler:
    # ...

    def forward(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> SequencePoolerOutput:
        # ...
        prompt_token_ids = pooling_metadata.prompt_token_ids
        # ...

Step 2: Verify the fix

Run the test again to verify that the fix works.

Verification

To verify that the fix worked, run the test_splade_pooler_matches_reference_formula test again. If the test passes, it means the fix was successful.

Extra Tips

Make sure to update the documentation and tests to reflect any changes made to the code. Additionally, consider adding a check to ensure that the pooling_metadata object has the required attributes to prevent similar issues in the future.