wehub-resource-sync
59a0a3844c
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
PR Test AMD / finish (push) Has been cancelled
PR Test NVIDIA / finish (push) Has been cancelled
PR Test NVIDIA ARM / finish (push) Has been cancelled
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
24 lines
731 B
Python
24 lines
731 B
Python
import torch
|
|
|
|
from tokenspeed.runtime.execution.context import ForwardContext
|
|
from tokenspeed.runtime.execution.forward_batch_info import ForwardMode
|
|
from tokenspeed.runtime.layers.logits_processor import LogitsMetadata
|
|
|
|
|
|
def test_logits_metadata_derives_basic_fields_from_forward_context():
|
|
gather_ids = torch.tensor([0], dtype=torch.int64)
|
|
ctx = ForwardContext(
|
|
attn_backend=None,
|
|
token_to_kv_pool=None,
|
|
bs=1,
|
|
num_extends=0,
|
|
input_num_tokens=1,
|
|
forward_mode=ForwardMode.DECODE,
|
|
gather_ids=gather_ids,
|
|
)
|
|
|
|
metadata = LogitsMetadata.from_forward_context(ctx)
|
|
|
|
assert metadata.forward_mode == ForwardMode.DECODE
|
|
assert metadata.gather_ids is gather_ids
|