[CI] Run the Clef pooler test on GPU (#61409)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Jiangyun Zhu
2026-10-13 14:42:53 +00:00
committed by GitHub
co-authored by Claude Opus 5.5
parent de8291e24f
commit 61d9c777bf
@@ -15,9 +15,11 @@ from vllm.model_executor.models.clef_joint_schema import (
JointSchemaHead,
JointSchemaPooler,
)
from vllm.platforms import current_platform
from vllm.pooling_params import PoolingParams
from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
DEVICE_TYPE = current_platform.device_type
HIDDEN, VOCAB = 64, 500
TYPES = [0, 1, 1, 2]
COUNTS = [2, 3, 1, 5]
@@ -55,7 +57,7 @@ def reference_logits(head, hidden, token_ids, embedding, question_spans, spans):
base
+ head.option_summary_norm(torch.stack(summaries))
+ head.global_projection(global_vector)
+ head.type_embedding(torch.tensor(TYPES))
+ head.type_embedding(torch.tensor(TYPES, device=hidden.device))
).unsqueeze(0)
for layer in head.layers:
fields = layer(fields, memory)
@@ -90,7 +92,8 @@ def pooler():
with torch.no_grad():
for p in head.parameters():
p.add_(torch.randn_like(p) * 0.2)
embedding = torch.randn(VOCAB, HIDDEN)
head.to(DEVICE_TYPE)
embedding = torch.randn(VOCAB, HIDDEN, device=DEVICE_TYPE)
config = SimpleNamespace(
scheduler_config=SimpleNamespace(
enable_chunked_prefill=False, async_scheduling=False
@@ -132,7 +135,7 @@ def test_pooler_matches_reference(pooler):
pooler, embedding = pooler
lengths = [120, 90, 100]
token_ids = [torch.randint(0, VOCAB, (n,)) for n in lengths]
hidden = torch.randn(sum(lengths), HIDDEN)
hidden = torch.randn(sum(lengths), HIDDEN, device=DEVICE_TYPE)
first, first_questions, first_spans = schema(lengths[0])
last, last_questions, last_spans = schema(lengths[2])
params = [
@@ -154,7 +157,7 @@ def test_pooler_matches_reference(pooler):
metadata.build_pooling_cursor(
torch.tensor(lengths).numpy(),
seq_lens_cpu=torch.tensor(lengths),
device=torch.device("cpu"),
device=hidden.device,
)
probs, empty, logits = pooler(hidden, metadata)