This commit is contained in:
Nick Hill
@@ -339,136 +339,6 @@ def test_dsv4_fast_topk(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="The DeepSeek V4 fast path is CUDA-only.",
|
||||
)
|
||||
def test_dsv4_fast_topk_padding_uint32_falls_back(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Padded rows need the -1 sentinel, which uint32 cannot represent: the
|
||||
router must skip the dsv4 fast path and still route the real rows."""
|
||||
torch.manual_seed(0)
|
||||
num_tokens = 17
|
||||
num_experts = 256
|
||||
hidden_states = torch.randn((num_tokens, 64), dtype=torch.float32, device="cuda")
|
||||
gating_output = torch.randn(
|
||||
(num_tokens, num_experts), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
|
||||
is_padding = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")
|
||||
is_padding[1::2] = True
|
||||
gating_output[is_padding] = float("nan")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm.model_executor.layers.fused_moe.router."
|
||||
"fused_topk_bias_router._get_padding_mask",
|
||||
lambda _: is_padding,
|
||||
)
|
||||
# uint32 + padding would trip dsv4_topk's signed-indices assertion; the
|
||||
# generic path must take over instead.
|
||||
topk_weights, topk_ids = fused_topk_bias(
|
||||
hidden_states=hidden_states,
|
||||
gating_output=gating_output,
|
||||
scoring_func="sqrtsoftplus",
|
||||
e_score_correction_bias=correction_bias,
|
||||
topk=6,
|
||||
renormalize=True,
|
||||
indices_type=torch.uint32,
|
||||
routed_scaling_factor=1.5,
|
||||
)
|
||||
|
||||
assert topk_ids.dtype == torch.uint32
|
||||
topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
|
||||
gating_output=gating_output[~is_padding],
|
||||
topk=6,
|
||||
renormalize=True,
|
||||
routed_scaling_factor=1.5,
|
||||
e_score_correction_bias=correction_bias,
|
||||
)
|
||||
# uint32 CUDA tensors do not support boolean-mask indexing; widen first.
|
||||
torch.testing.assert_close(
|
||||
topk_ids.to(torch.int64)[~is_padding],
|
||||
topk_ids_ref.to(torch.int64),
|
||||
atol=0,
|
||||
rtol=0,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="The DeepSeek V4 fast path is CUDA-only.",
|
||||
)
|
||||
def test_dsv4_fast_topk_padding(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Verify the DSV4 fast path removes graph-padding rows from routing."""
|
||||
torch.manual_seed(0)
|
||||
num_tokens = 17
|
||||
num_experts = 256
|
||||
hidden_states = torch.randn((num_tokens, 64), dtype=torch.float32, device="cuda")
|
||||
gating_output = torch.randn(
|
||||
(num_tokens, num_experts), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
|
||||
is_padding = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")
|
||||
is_padding[1::2] = True
|
||||
gating_output[is_padding] = float("nan")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm.model_executor.layers.fused_moe.router."
|
||||
"fused_topk_bias_router._get_padding_mask",
|
||||
lambda _: is_padding,
|
||||
)
|
||||
topk_weights, topk_ids = fused_topk_bias(
|
||||
hidden_states=hidden_states,
|
||||
gating_output=gating_output,
|
||||
scoring_func="sqrtsoftplus",
|
||||
e_score_correction_bias=correction_bias,
|
||||
topk=6,
|
||||
renormalize=True,
|
||||
routed_scaling_factor=1.5,
|
||||
)
|
||||
|
||||
assert torch.equal(topk_ids[is_padding], torch.full_like(topk_ids[is_padding], -1))
|
||||
assert torch.equal(
|
||||
topk_weights[is_padding], torch.zeros_like(topk_weights[is_padding])
|
||||
)
|
||||
|
||||
topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
|
||||
gating_output=gating_output[~is_padding],
|
||||
topk=6,
|
||||
renormalize=True,
|
||||
routed_scaling_factor=1.5,
|
||||
e_score_correction_bias=correction_bias,
|
||||
)
|
||||
torch.testing.assert_close(topk_ids[~is_padding], topk_ids_ref, atol=0, rtol=0)
|
||||
torch.testing.assert_close(
|
||||
topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5
|
||||
)
|
||||
|
||||
# The mask buffer is persistent under CUDA graph replay, but its contents
|
||||
# change with every batch. Verify that the kernel reads those contents at
|
||||
# runtime rather than specializing on the mask captured above.
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
graph_weights, graph_ids = dsv4_topk(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
torch.int32,
|
||||
1.5,
|
||||
is_padding=is_padding,
|
||||
)
|
||||
|
||||
is_padding.logical_not_()
|
||||
graph.replay()
|
||||
assert torch.equal(
|
||||
graph_ids[is_padding], torch.full_like(graph_ids[is_padding], -1)
|
||||
)
|
||||
assert torch.equal(
|
||||
graph_weights[is_padding], torch.zeros_like(graph_weights[is_padding])
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="This test is skipped on non-CUDA platform.",
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
@@ -9,9 +8,7 @@ import pytest
|
||||
import torch
|
||||
|
||||
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
|
||||
from vllm.v1.worker.gpu.spec_decode.dflash import speculator as dflash_speculator
|
||||
from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
|
||||
DFlashSpeculator,
|
||||
prepare_dflash_inputs,
|
||||
)
|
||||
|
||||
@@ -36,7 +33,6 @@ def _run_prepare(
|
||||
input_buffers = SimpleNamespace(
|
||||
input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device),
|
||||
positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device),
|
||||
is_padding=torch.zeros(max_num_tokens, dtype=torch.bool, device=device),
|
||||
query_start_loc=torch.full(
|
||||
(max_num_reqs + 1,), -1, dtype=torch.int32, device=device
|
||||
),
|
||||
@@ -144,114 +140,6 @@ def test_prepare_dflash_inputs_excludes_rejected_context_suffix():
|
||||
assert out.temperature[2].item() == 1.0
|
||||
assert out.seeds[2].item() == 17
|
||||
|
||||
assert not out.input_buffers.is_padding[:3].any()
|
||||
assert out.input_buffers.is_padding[3:].all()
|
||||
assert out.input_buffers.input_ids[3:].cpu().tolist() == [0] * 13
|
||||
assert out.input_buffers.positions[3:].cpu().tolist() == [0] * 13
|
||||
|
||||
|
||||
def test_prepare_dflash_inputs_compacts_noncontiguous_request_slots():
|
||||
device = torch.device("cuda")
|
||||
max_num_reqs = 4
|
||||
max_num_tokens = 16
|
||||
num_speculative_steps = 3
|
||||
input_buffers = SimpleNamespace(
|
||||
input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device),
|
||||
positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device),
|
||||
is_padding=torch.zeros(max_num_tokens, dtype=torch.bool, device=device),
|
||||
query_start_loc=torch.full(
|
||||
(max_num_reqs + 1,), -1, dtype=torch.int32, device=device
|
||||
),
|
||||
seq_lens=torch.full((max_num_reqs,), -1, dtype=torch.int32, device=device),
|
||||
)
|
||||
input_batch = SimpleNamespace(
|
||||
num_reqs=2,
|
||||
num_scheduled_tokens=np.array([4, 4], dtype=np.int32),
|
||||
positions=torch.tensor(
|
||||
[10, 11, 12, 13, 20, 21, 22, 23],
|
||||
dtype=torch.int64,
|
||||
device=device,
|
||||
),
|
||||
query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32, device=device),
|
||||
# Active batch rows are compact, while request state remains in slots 3 and 1.
|
||||
idx_mapping=torch.tensor([3, 1], dtype=torch.int32, device=device),
|
||||
)
|
||||
query_slot_mapping = torch.full(
|
||||
(max_num_tokens,), -2, dtype=torch.int64, device=device
|
||||
)
|
||||
context_positions = torch.full(
|
||||
(max_num_tokens,), -1, dtype=torch.int64, device=device
|
||||
)
|
||||
context_slot_mapping = torch.full(
|
||||
(max_num_tokens,), -2, dtype=torch.int64, device=device
|
||||
)
|
||||
sample_indices = torch.full(
|
||||
(max_num_reqs * num_speculative_steps,),
|
||||
-1,
|
||||
dtype=torch.int64,
|
||||
device=device,
|
||||
)
|
||||
sample_pos = torch.full_like(sample_indices, -1)
|
||||
sample_idx_mapping = torch.full(
|
||||
sample_indices.shape, -1, dtype=torch.int32, device=device
|
||||
)
|
||||
temperature = torch.zeros(max_num_reqs, dtype=torch.float32, device=device)
|
||||
seeds = torch.zeros(max_num_reqs, dtype=torch.int64, device=device)
|
||||
input_temperature = torch.tensor(
|
||||
[0.0, 0.5, 0.0, 1.0], dtype=torch.float32, device=device
|
||||
)
|
||||
input_seeds = torch.tensor([0, 11, 0, 33], dtype=torch.int64, device=device)
|
||||
last_sampled = torch.tensor([0, 77, 0, 99], dtype=torch.int64, device=device)
|
||||
next_prefill_tokens = torch.zeros_like(last_sampled)
|
||||
block_table = torch.tensor(
|
||||
[[0, 0, 7, 8, 9, 10, 11, 12], [0, 0, 13, 14, 15, 16, 17, 18]],
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
prepare_dflash_inputs(
|
||||
input_buffers,
|
||||
query_slot_mapping,
|
||||
context_positions,
|
||||
context_slot_mapping,
|
||||
sample_indices,
|
||||
sample_pos,
|
||||
sample_idx_mapping,
|
||||
temperature,
|
||||
seeds,
|
||||
input_batch,
|
||||
torch.tensor([1, 1], dtype=torch.int32, device=device),
|
||||
torch.tensor([2, 1], dtype=torch.int32, device=device),
|
||||
last_sampled,
|
||||
next_prefill_tokens,
|
||||
input_temperature,
|
||||
input_seeds,
|
||||
block_table,
|
||||
4,
|
||||
0,
|
||||
1,
|
||||
1,
|
||||
123,
|
||||
num_speculative_steps,
|
||||
num_speculative_steps,
|
||||
max_num_reqs,
|
||||
max_num_tokens,
|
||||
128,
|
||||
sample_from_anchor=True,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# Query rows follow compact batch order, but every persistent state lookup
|
||||
# follows idx_mapping instead of accidentally using the compact row index.
|
||||
assert input_buffers.input_ids[:6].cpu().tolist() == [99, 123, 123, 77, 123, 123]
|
||||
assert input_buffers.positions[:6].cpu().tolist() == [12, 13, 14, 23, 24, 25]
|
||||
assert sample_indices[:6].cpu().tolist() == [0, 1, 2, 3, 4, 5]
|
||||
assert sample_idx_mapping[:6].cpu().tolist() == [3, 3, 3, 1, 1, 1]
|
||||
assert temperature.cpu().tolist() == [0.0, 0.5, 0.0, 1.0]
|
||||
assert seeds.cpu().tolist() == [0, 11, 0, 33]
|
||||
assert not input_buffers.is_padding[:6].any()
|
||||
assert input_buffers.is_padding[6:].all()
|
||||
|
||||
|
||||
def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp():
|
||||
out = _run_prepare(
|
||||
@@ -286,40 +174,3 @@ def test_prepare_dflash_inputs_never_writes_the_null_block():
|
||||
PAD_SLOT_ID,
|
||||
PAD_SLOT_ID,
|
||||
]
|
||||
|
||||
|
||||
def test_dflash_forward_context_receives_draft_padding_mask(monkeypatch):
|
||||
device = torch.device("cuda")
|
||||
input_buffers = SimpleNamespace(
|
||||
input_ids=torch.tensor([11, 12, 0, 0], dtype=torch.int32, device=device),
|
||||
positions=torch.tensor([7, 8, 0, 0], dtype=torch.int64, device=device),
|
||||
is_padding=torch.tensor([False, False, True, True], device=device),
|
||||
)
|
||||
observed = None
|
||||
|
||||
@contextmanager
|
||||
def fake_set_forward_context(*args, **kwargs):
|
||||
nonlocal observed
|
||||
observed = kwargs["is_padding"].clone()
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(
|
||||
dflash_speculator, "set_forward_context", fake_set_forward_context
|
||||
)
|
||||
speculator = SimpleNamespace(
|
||||
input_buffers=input_buffers,
|
||||
vllm_config=SimpleNamespace(),
|
||||
model=lambda **kwargs: kwargs["input_ids"],
|
||||
)
|
||||
|
||||
result = DFlashSpeculator._run_model(
|
||||
speculator,
|
||||
num_tokens=4,
|
||||
attn_metadata=None,
|
||||
slot_mappings=None,
|
||||
num_tokens_across_dp=None,
|
||||
)
|
||||
|
||||
assert result.tolist() == [11, 12, 0, 0]
|
||||
assert observed is not None
|
||||
assert observed.tolist() == [False, False, True, True]
|
||||
|
||||
@@ -5,35 +5,16 @@
|
||||
from unittest.mock import Mock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.v1.worker.gpu import model_runner, pp_utils
|
||||
from vllm.v1.worker.gpu import pp_utils
|
||||
|
||||
|
||||
def _cuda_handler(max_sample_len=6):
|
||||
handler = object.__new__(pp_utils.PPHandler)
|
||||
handler.is_last_rank = True
|
||||
handler.disabled = False
|
||||
handler.max_sample_len = max_sample_len
|
||||
handler.last_rank = 1
|
||||
handler.broadcast_group = Mock()
|
||||
handler.device = torch.device("cuda")
|
||||
handler.main_stream = torch.cuda.current_stream()
|
||||
handler.broadcast_stream = torch.cuda.Stream()
|
||||
return handler
|
||||
|
||||
|
||||
def _batch(num_computed, prefill_len, num_scheduled, idx_mapping=None):
|
||||
num_reqs = len(num_computed)
|
||||
if idx_mapping is None:
|
||||
idx_mapping = list(range(num_reqs))
|
||||
def _batch(num_computed, prefill_len, num_scheduled):
|
||||
return Mock(
|
||||
num_reqs=num_reqs,
|
||||
num_reqs=len(num_computed),
|
||||
num_computed_tokens_np=np.array(num_computed, dtype=np.int32),
|
||||
prefill_len_np=np.array(prefill_len, dtype=np.int32),
|
||||
num_scheduled_tokens=np.array(num_scheduled, dtype=np.int32),
|
||||
idx_mapping=torch.tensor(idx_mapping, dtype=torch.int64),
|
||||
)
|
||||
|
||||
|
||||
@@ -100,130 +81,3 @@ def test_decode_row_ahead_of_a_prefill_chunk():
|
||||
|
||||
assert mask is not None
|
||||
assert mask.tolist() == [True, False]
|
||||
|
||||
|
||||
def test_disabled_handler_skips_broadcast_and_receive(monkeypatch):
|
||||
"""While disabled (warmup), neither side enqueues a broadcast op."""
|
||||
sent = []
|
||||
monkeypatch.setattr(
|
||||
pp_utils.torch.distributed,
|
||||
"broadcast",
|
||||
lambda *args, **kwargs: sent.append((args, kwargs)),
|
||||
)
|
||||
|
||||
handler = object.__new__(pp_utils.PPHandler)
|
||||
handler.set_disabled(True)
|
||||
|
||||
handler.is_last_rank = False
|
||||
assert handler.receive(Mock()) is False
|
||||
|
||||
handler.is_last_rank = True
|
||||
assert handler.broadcast(Mock(), Mock(), Mock(), Mock()) is None
|
||||
|
||||
assert sent == []
|
||||
|
||||
handler.set_disabled(False)
|
||||
assert handler.disabled is False
|
||||
|
||||
|
||||
def test_alloc_combined_keeps_unbind_views_16_byte_aligned():
|
||||
"""Triton specializes on pointer alignment: an unaligned `num_rejected`
|
||||
would compile a second `_post_update_kernel` variant at serving time,
|
||||
where the in-flight broadcast NCCL kernel can block the module load."""
|
||||
for num_reqs in range(1, 9):
|
||||
combined = pp_utils._alloc_combined(num_reqs, torch.device("cpu"))
|
||||
num_sampled, num_rejected = combined.unbind(dim=0)
|
||||
assert num_sampled.data_ptr() % 16 == 0
|
||||
assert num_rejected.data_ptr() % 16 == 0
|
||||
assert combined.shape[1] >= num_reqs
|
||||
|
||||
|
||||
def test_warmup_pp_decode_update_matches_serving_specialization(monkeypatch):
|
||||
"""The warmup launch must hit the same triton specialization as serving.
|
||||
|
||||
A mismatch means the first real ``update_pp_decode_requests`` recompiles
|
||||
mid-serving, where the in-flight broadcast NCCL kernel blocks the CUDA
|
||||
module load and deadlocks the pipeline.
|
||||
"""
|
||||
calls = []
|
||||
monkeypatch.setattr(model_runner, "post_update", lambda *args: calls.append(args))
|
||||
|
||||
runner = object.__new__(model_runner.GPUModelRunner)
|
||||
runner.device = torch.device("cpu")
|
||||
runner.pp_handler = Mock(max_sample_len=3)
|
||||
runner.req_states = Mock()
|
||||
|
||||
runner.warmup_pp_decode_update()
|
||||
|
||||
assert len(calls) == 1
|
||||
args = calls[0]
|
||||
idx_mapping, _, _, output_bin_counts = args[:4]
|
||||
sampled_tokens, num_sampled, num_rejected, query_start_loc = args[4:8]
|
||||
broadcast_drafts, draft_tokens_out = args[10:12]
|
||||
assert idx_mapping.tolist() == [-1] and idx_mapping.dtype == torch.int64
|
||||
assert output_bin_counts is None
|
||||
assert query_start_loc is None
|
||||
assert sampled_tokens.shape == (1, 3) and sampled_tokens.dtype == torch.int64
|
||||
assert num_sampled.dtype == torch.int32
|
||||
assert num_rejected.dtype == torch.int32
|
||||
# Spec-enabled PP handlers receive drafts over the broadcast; the warmup
|
||||
# must compile that specialization (non-None draft pointers) too.
|
||||
assert broadcast_drafts.shape == (1, 2) and broadcast_drafts.dtype == torch.int64
|
||||
assert draft_tokens_out is runner.req_states.draft_tokens
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA stream")
|
||||
def test_broadcast_pads_plain_sampler_rows_to_max_sample_len(monkeypatch):
|
||||
"""The wire shape must not depend on whether the batch carried drafts:
|
||||
the receiver always allocates [num_reqs, max_sample_len], and a NCCL
|
||||
broadcast with mismatched counts hangs the receiver."""
|
||||
sent = []
|
||||
monkeypatch.setattr(
|
||||
pp_utils.torch.distributed,
|
||||
"broadcast",
|
||||
lambda tensor, **kw: sent.append(tensor),
|
||||
)
|
||||
handler = _cuda_handler()
|
||||
batch = _batch(num_computed=[10], prefill_len=[8], num_scheduled=[1])
|
||||
|
||||
handler.broadcast(
|
||||
torch.zeros(1, 1, dtype=torch.int64, device="cuda"), # plain sampler
|
||||
torch.ones(1, dtype=torch.int32, device="cuda"),
|
||||
torch.zeros(1, dtype=torch.int32, device="cuda"),
|
||||
batch,
|
||||
)
|
||||
|
||||
assert sent[0].shape == (1, 6)
|
||||
assert sent[1].shape == (2, 4)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA stream")
|
||||
def test_broadcast_drafts_gathers_fresh_rows_from_the_table(monkeypatch):
|
||||
sent = []
|
||||
monkeypatch.setattr(
|
||||
pp_utils.torch.distributed,
|
||||
"broadcast",
|
||||
lambda tensor, **kw: sent.append(tensor),
|
||||
)
|
||||
handler = _cuda_handler()
|
||||
# The batch's single row maps to request-state row 2.
|
||||
batch = _batch(num_computed=[10], prefill_len=[8], num_scheduled=[1])
|
||||
batch.idx_mapping = torch.tensor([2], dtype=torch.int64, device="cuda")
|
||||
table = torch.arange(20, dtype=torch.int64, device="cuda").view(4, 5)
|
||||
|
||||
handler.broadcast_drafts(table, batch)
|
||||
|
||||
assert sent[0].shape == (1, 5)
|
||||
# The payload is a gather into a fresh tensor: propose() overwrites its
|
||||
# persistent buffer on the next step, possibly before this send completes.
|
||||
assert sent[0].data_ptr() != table.data_ptr()
|
||||
assert sent[0].cpu().tolist() == [table[2].cpu().tolist()]
|
||||
|
||||
# An all-prefill batch sends nothing (receive() enqueues nothing either).
|
||||
sent.clear()
|
||||
prefill_batch = _batch(num_computed=[0], prefill_len=[4096], num_scheduled=[448])
|
||||
prefill_batch.idx_mapping = prefill_batch.idx_mapping.cuda()
|
||||
handler.broadcast_drafts(table, prefill_batch)
|
||||
assert sent == []
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
@@ -41,7 +41,6 @@ if current_platform.is_cuda():
|
||||
def _dsv4_topk_kernel(
|
||||
gating_output_ptr,
|
||||
correction_bias_ptr,
|
||||
is_padding_ptr,
|
||||
topk_weights_ptr,
|
||||
topk_ids_ptr,
|
||||
routed_scaling_factor,
|
||||
@@ -50,7 +49,6 @@ if current_platform.is_cuda():
|
||||
image_sentinel_lo,
|
||||
NUM_EXPERTS: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
HAS_PADDING: tl.constexpr,
|
||||
HAS_VL: tl.constexpr,
|
||||
launch_pdl: tl.constexpr,
|
||||
):
|
||||
@@ -108,11 +106,6 @@ if current_platform.is_cuda():
|
||||
output_mask = topk_offsets < 6
|
||||
output_offsets = row * 6 + topk_offsets
|
||||
|
||||
if HAS_PADDING:
|
||||
is_padding = tl.load(is_padding_ptr + row)
|
||||
selected_weights = tl.where(is_padding, 0.0, selected_weights)
|
||||
selected_ids = tl.where(is_padding, -1, selected_ids)
|
||||
|
||||
if launch_pdl:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
@@ -125,20 +118,11 @@ def dsv4_topk(
|
||||
correction_bias: torch.Tensor,
|
||||
indices_dtype: torch.dtype,
|
||||
routed_scaling_factor: float,
|
||||
is_padding: torch.Tensor | None = None,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
bias_vl: torch.Tensor | None = None,
|
||||
image_sentinel_lo: int = 0,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
num_tokens, num_experts = gating_output.shape
|
||||
if is_padding is not None:
|
||||
assert is_padding.dtype == torch.bool
|
||||
assert is_padding.shape == (num_tokens,)
|
||||
assert is_padding.device == gating_output.device
|
||||
assert is_padding.is_contiguous()
|
||||
assert indices_dtype in (torch.int32, torch.int64), (
|
||||
"Padding requires a signed indices dtype for the -1 sentinel."
|
||||
)
|
||||
has_vl = bias_vl is not None and image_sentinel_lo > 0
|
||||
if bias_vl is not None:
|
||||
assert input_ids is not None, "bias_vl routing requires input_ids"
|
||||
@@ -152,7 +136,6 @@ def dsv4_topk(
|
||||
_dsv4_topk_kernel[(num_tokens,)](
|
||||
gating_output,
|
||||
correction_bias,
|
||||
is_padding,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
routed_scaling_factor,
|
||||
@@ -161,7 +144,6 @@ def dsv4_topk(
|
||||
image_sentinel_lo,
|
||||
NUM_EXPERTS=num_experts,
|
||||
BLOCK_N=triton.next_power_of_2(num_experts),
|
||||
HAS_PADDING=is_padding is not None,
|
||||
HAS_VL=has_vl,
|
||||
num_warps=1,
|
||||
launch_pdl=current_platform.is_arch_support_pdl(),
|
||||
|
||||
@@ -148,13 +148,9 @@ def fused_topk_bias(
|
||||
)
|
||||
|
||||
output_indices_dtype = torch.int32 if indices_type is None else indices_type
|
||||
padding_mask = _get_padding_mask(gating_output.shape[0])
|
||||
if (
|
||||
scoring_func == "sqrtsoftplus"
|
||||
and hash_indices_table is None
|
||||
# dsv4_topk marks padded rows with a -1 sentinel, which unsigned
|
||||
# indices cannot represent; keep the generic path for that pair.
|
||||
and (padding_mask is None or output_indices_dtype != torch.uint32)
|
||||
and can_use_dsv4_topk(
|
||||
gating_output,
|
||||
e_score_correction_bias,
|
||||
@@ -169,7 +165,6 @@ def fused_topk_bias(
|
||||
e_score_correction_bias,
|
||||
output_indices_dtype,
|
||||
routed_scaling_factor,
|
||||
is_padding=padding_mask,
|
||||
input_ids=input_tokens,
|
||||
bias_vl=bias_vl,
|
||||
image_sentinel_lo=image_sentinel_lo,
|
||||
|
||||
@@ -555,11 +555,6 @@ def _post_update_kernel(
|
||||
all_token_ids_ptr,
|
||||
all_token_ids_stride,
|
||||
total_len_ptr,
|
||||
broadcast_drafts_ptr,
|
||||
broadcast_drafts_stride,
|
||||
draft_tokens_ptr,
|
||||
draft_tokens_stride,
|
||||
num_spec,
|
||||
):
|
||||
req_id = tl.program_id(0)
|
||||
req_state_idx = tl.load(idx_mapping_ptr + req_id)
|
||||
@@ -567,17 +562,6 @@ def _post_update_kernel(
|
||||
# Filter rows with negative index entries.
|
||||
return
|
||||
|
||||
if broadcast_drafts_ptr is not None:
|
||||
# PP path: adopt the draft tokens proposed by the last rank's
|
||||
# speculator so the next verification step embeds the real drafts.
|
||||
for i in range(num_spec):
|
||||
token_id = tl.load(
|
||||
broadcast_drafts_ptr + req_id * broadcast_drafts_stride + i
|
||||
)
|
||||
tl.store(
|
||||
draft_tokens_ptr + req_state_idx * draft_tokens_stride + i, token_id
|
||||
)
|
||||
|
||||
total_len = tl.load(total_len_ptr + req_state_idx)
|
||||
num_sampled = tl.load(num_sampled_ptr + req_id)
|
||||
if num_sampled > 0:
|
||||
@@ -638,11 +622,6 @@ def post_update(
|
||||
all_token_ids: torch.Tensor,
|
||||
# [max_num_reqs]
|
||||
total_len: torch.Tensor,
|
||||
# [num_reqs, num_spec]; drafts broadcast from the last PP rank. Only
|
||||
# passed on non-last PP ranks, which never run the speculator.
|
||||
broadcast_drafts: torch.Tensor | None = None,
|
||||
# [max_num_reqs, num_spec]
|
||||
draft_tokens_out: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
num_reqs = idx_mapping.shape[0]
|
||||
_post_update_kernel[(num_reqs,)](
|
||||
@@ -659,11 +638,6 @@ def post_update(
|
||||
all_token_ids,
|
||||
all_token_ids.stride(0),
|
||||
total_len,
|
||||
broadcast_drafts,
|
||||
broadcast_drafts.stride(0) if broadcast_drafts is not None else 0,
|
||||
draft_tokens_out,
|
||||
draft_tokens_out.stride(0) if draft_tokens_out is not None else 0,
|
||||
broadcast_drafts.shape[1] if broadcast_drafts is not None else 0,
|
||||
num_warps=1,
|
||||
)
|
||||
|
||||
|
||||
@@ -862,7 +862,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
||||
# dummy run the eagle speculator's propose to ensure DP/EP sync.
|
||||
if self.speculator is not None:
|
||||
assert self.sampler is not None
|
||||
assert hidden_states is not None
|
||||
self.step_timing.drafter_start()
|
||||
mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
|
||||
if self.speculator.supports_mm_inputs:
|
||||
@@ -878,8 +877,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
||||
spec_hidden_states = hidden_states
|
||||
if hasattr(self.model, "get_mtp_target_hidden_states"):
|
||||
pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
|
||||
if pre_hc_hidden_states is not None:
|
||||
spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]
|
||||
spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr]
|
||||
if isinstance(self.sampler, GPUWatermarkSampler):
|
||||
self.speculator.prepare_watermarking(
|
||||
self.sampler._get_contexts(input_batch.idx_mapping),
|
||||
@@ -1111,44 +1109,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
||||
if outputs is not None:
|
||||
self.postprocess_sampled(**outputs)
|
||||
|
||||
def warmup_pp_decode_update(self) -> None:
|
||||
"""JIT-compile the kernel behind ``update_pp_decode_requests``.
|
||||
|
||||
That path only runs on real steps, so the warmup steps never reach it
|
||||
on non-last PP ranks. Its first triton compile must not happen
|
||||
mid-serving: the in-flight sampled-token broadcast keeps a NCCL kernel
|
||||
spinning on this device, which blocks the CUDA module load and
|
||||
deadlocks the pipeline. An all -1 idx_mapping makes this a no-op.
|
||||
The freshly allocated int32 tensors are 16-byte aligned, matching the
|
||||
padded views `PPHandler` produces at serving time (triton specializes
|
||||
on pointer alignment).
|
||||
"""
|
||||
assert self.pp_handler is not None
|
||||
num_spec = self.pp_handler.max_sample_len - 1
|
||||
broadcast_drafts = (
|
||||
torch.zeros((1, num_spec), dtype=torch.int64, device=self.device)
|
||||
if num_spec > 0
|
||||
else None
|
||||
)
|
||||
post_update(
|
||||
torch.full((1,), -1, dtype=torch.int64, device=self.device),
|
||||
self.req_states.num_computed_tokens.gpu,
|
||||
self.req_states.last_sampled_tokens,
|
||||
None,
|
||||
torch.zeros(
|
||||
(1, self.pp_handler.max_sample_len),
|
||||
dtype=torch.int64,
|
||||
device=self.device,
|
||||
),
|
||||
torch.zeros(1, dtype=torch.int32, device=self.device),
|
||||
torch.zeros(1, dtype=torch.int32, device=self.device),
|
||||
None,
|
||||
self.req_states.all_token_ids.gpu,
|
||||
self.req_states.total_len.gpu,
|
||||
broadcast_drafts,
|
||||
self.req_states.draft_tokens if broadcast_drafts is not None else None,
|
||||
)
|
||||
|
||||
def add_requests(self, scheduler_output: SchedulerOutput) -> None:
|
||||
for new_req_data in scheduler_output.scheduled_new_reqs:
|
||||
assert new_req_data.prefill_token_ids is not None
|
||||
@@ -1630,7 +1590,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
||||
num_sampled: torch.Tensor,
|
||||
num_rejected: torch.Tensor,
|
||||
query_start_loc: torch.Tensor | None = None,
|
||||
broadcast_drafts: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
# Update the number of computed tokens.
|
||||
output_bin_counts = None
|
||||
@@ -1648,8 +1607,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
||||
query_start_loc,
|
||||
self.req_states.all_token_ids.gpu,
|
||||
self.req_states.total_len.gpu,
|
||||
broadcast_drafts,
|
||||
self.req_states.draft_tokens if broadcast_drafts is not None else None,
|
||||
)
|
||||
|
||||
self.model_state.postprocess_state(
|
||||
@@ -2170,10 +2127,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
||||
spec_hidden_states = draft_hidden_states
|
||||
if hasattr(self.model, "get_mtp_target_hidden_states"):
|
||||
pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
|
||||
if pre_hc_hidden_states is not None:
|
||||
spec_hidden_states = pre_hc_hidden_states[
|
||||
: draft_hidden_states.size(0)
|
||||
]
|
||||
spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]
|
||||
if isinstance(self.sampler, GPUWatermarkSampler):
|
||||
self.speculator.prepare_watermarking(
|
||||
self.sampler._get_contexts(input_batch.idx_mapping),
|
||||
@@ -2195,16 +2149,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
||||
dp_sync=dp_sync,
|
||||
mm_inputs=mm_inputs,
|
||||
)
|
||||
if draft_tokens is not None:
|
||||
self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
|
||||
if self.pp_handler is not None:
|
||||
# Earlier stages never run the speculator; ship the
|
||||
# drafts so their next verification step embeds the real
|
||||
# draft tokens instead of stale buffer contents.
|
||||
self.pp_handler.broadcast_drafts(
|
||||
self.req_states.draft_tokens, input_batch
|
||||
)
|
||||
if draft_tokens is not None and self.adaptive_verification is not None:
|
||||
self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
|
||||
if self.adaptive_verification is not None:
|
||||
self.adaptive_verification.record_confidences(
|
||||
self.speculator.draft_token_confidence_probs, input_batch
|
||||
)
|
||||
@@ -2216,11 +2162,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
||||
input_batch,
|
||||
self.req_states.draft_tokens[input_batch.idx_mapping],
|
||||
)
|
||||
if self.pp_handler is not None and self.speculator is None:
|
||||
# When a speculator ran, the propose() path above already
|
||||
# broadcast the fresh drafts. Broadcasting here as well would
|
||||
# double-post on the pp_broadcast group and misalign the
|
||||
# recv FIFO on earlier stages, hanging the pipeline.
|
||||
if self.pp_handler is not None:
|
||||
self.pp_handler.broadcast_drafts(
|
||||
self.req_states.draft_tokens, input_batch
|
||||
)
|
||||
|
||||
@@ -46,18 +46,6 @@ def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None:
|
||||
return produces_sample if produces_sample.any() else None
|
||||
|
||||
|
||||
def _alloc_combined(num_reqs: int, device: torch.device) -> torch.Tensor:
|
||||
"""Allocate the (2, N) int32 buffer broadcast alongside sampled tokens.
|
||||
|
||||
The inner dim is padded to a multiple of 4 so that both `unbind` views
|
||||
stay 16-byte aligned for any `num_reqs`: triton specializes on pointer
|
||||
alignment, and a misaligned `num_rejected` would JIT-compile a second
|
||||
`_post_update_kernel` variant at serving time. The padding is broadcast
|
||||
but never read. Sender and receiver must both use this allocation.
|
||||
"""
|
||||
return torch.empty(2, -(-num_reqs // 4) * 4, dtype=torch.int32, device=device)
|
||||
|
||||
|
||||
class PPHandler:
|
||||
"""Runs the PP sampled-token broadcast/recv on a side stream so the
|
||||
default stream isn't gated by the matching peer call. Step T's recv is
|
||||
@@ -98,14 +86,6 @@ class PPHandler:
|
||||
)
|
||||
self.aux_hidden_state_relay_keys: tuple[str, ...] = ()
|
||||
|
||||
# Warmup steps run the pipeline with synthetic batches whose outputs are
|
||||
# discarded; the sampled-token broadcast is disabled there so its
|
||||
# side-stream NCCL ops cannot overlap the next step's activation p2p.
|
||||
self.disabled = False
|
||||
|
||||
def set_disabled(self, disabled: bool) -> None:
|
||||
self.disabled = disabled
|
||||
|
||||
def on_req_idx_freed(self, req_idx: int) -> None:
|
||||
self.req_idx_gen_np[req_idx] += 1
|
||||
|
||||
@@ -177,15 +157,29 @@ class PPHandler:
|
||||
num_sampled=slot.num_sampled,
|
||||
num_rejected=slot.num_rejected,
|
||||
idx_mapping=idx_mapping,
|
||||
broadcast_drafts=slot.draft_tokens,
|
||||
)
|
||||
|
||||
def broadcast_drafts(
|
||||
self, draft_tokens: torch.Tensor, input_batch: InputBatch
|
||||
) -> None:
|
||||
"""Broadcast draft proposals so non-last ranks can embed real token ids."""
|
||||
assert self.is_last_rank
|
||||
if compute_need_sampled_mask(input_batch) is None:
|
||||
return
|
||||
with torch.cuda.stream(self.broadcast_stream):
|
||||
self.broadcast_stream.wait_stream(self.main_stream)
|
||||
send = draft_tokens[input_batch.idx_mapping].contiguous()
|
||||
# Must record the idx_mapping tensor since it was allocated
|
||||
# on the main stream.
|
||||
input_batch.idx_mapping.record_stream(self.broadcast_stream)
|
||||
torch.distributed.broadcast(
|
||||
send, src=self.last_rank, group=self.broadcast_group
|
||||
)
|
||||
|
||||
def receive(self, input_batch: InputBatch) -> bool:
|
||||
"""Returns True iff sampled tokens need to be gathered from *all*
|
||||
requests in the batch."""
|
||||
assert not self.is_last_rank
|
||||
if self.disabled:
|
||||
return False
|
||||
need_sampled_mask = compute_need_sampled_mask(input_batch)
|
||||
if need_sampled_mask is None:
|
||||
# Leave this step's reserved slot as None.
|
||||
@@ -201,7 +195,7 @@ class PPHandler:
|
||||
sampled_tokens = torch.empty(
|
||||
num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device
|
||||
)
|
||||
combined = _alloc_combined(num_reqs, self.device)
|
||||
combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device)
|
||||
torch.distributed.broadcast(
|
||||
sampled_tokens, src=self.last_rank, group=self.broadcast_group
|
||||
)
|
||||
@@ -219,7 +213,6 @@ class PPHandler:
|
||||
torch.distributed.broadcast(
|
||||
draft_tokens, src=self.last_rank, group=self.broadcast_group
|
||||
)
|
||||
|
||||
event = self.broadcast_stream.record_event()
|
||||
num_sampled, num_rejected = combined.unbind(dim=0)
|
||||
# Must record_stream since these were allocated on broadcast stream but
|
||||
@@ -249,16 +242,11 @@ class PPHandler:
|
||||
input_batch: InputBatch,
|
||||
) -> None:
|
||||
assert self.is_last_rank
|
||||
if self.disabled:
|
||||
return
|
||||
mask = compute_need_sampled_mask(input_batch)
|
||||
if mask is None:
|
||||
if compute_need_sampled_mask(input_batch) is None:
|
||||
# No request needs sampled outputs for a subsequent decode step.
|
||||
return
|
||||
|
||||
assert sampled_token_ids.dtype == torch.int64
|
||||
assert num_sampled.dtype == torch.int32
|
||||
assert num_rejected.dtype == torch.int32
|
||||
|
||||
if current_platform.is_xpu():
|
||||
self.main_stream.synchronize()
|
||||
@@ -274,42 +262,9 @@ class PPHandler:
|
||||
src=self.last_rank,
|
||||
group=self.broadcast_group,
|
||||
)
|
||||
combined = _alloc_combined(num_sampled.shape[0], self.device)
|
||||
combined[0, : num_sampled.shape[0]] = num_sampled
|
||||
combined[1, : num_sampled.shape[0]] = num_rejected
|
||||
combined = torch.stack((num_sampled, num_rejected), dim=0)
|
||||
torch.distributed.broadcast(
|
||||
combined, src=self.last_rank, group=self.broadcast_group
|
||||
)
|
||||
for tensor in (sampled_token_ids, num_sampled, num_rejected):
|
||||
tensor.record_stream(self.broadcast_stream)
|
||||
|
||||
def broadcast_drafts(
|
||||
self, draft_token_table: torch.Tensor, input_batch: InputBatch
|
||||
) -> None:
|
||||
"""Broadcast the speculator's freshly proposed draft tokens.
|
||||
|
||||
Runs after propose() on the last rank; the send is stream-ordered
|
||||
after broadcast()'s sends, matching receive()'s enqueue order. The
|
||||
payload is gathered from the runner's draft table (just updated from
|
||||
propose()'s output) into a fresh compact tensor: the speculator
|
||||
overwrites its own persistent buffer on the next step, possibly
|
||||
before this async send completes.
|
||||
"""
|
||||
assert self.is_last_rank
|
||||
if self.disabled or self.max_sample_len == 1:
|
||||
return
|
||||
if compute_need_sampled_mask(input_batch) is None:
|
||||
return
|
||||
assert draft_token_table.dtype == torch.int64
|
||||
assert draft_token_table.shape[1] == self.max_sample_len - 1
|
||||
|
||||
# Gather on the main stream so the payload is ordered after this
|
||||
# step's table update and before any later one.
|
||||
drafts = draft_token_table[input_batch.idx_mapping]
|
||||
assert drafts.shape[0] == input_batch.num_reqs
|
||||
with torch.cuda.stream(self.broadcast_stream):
|
||||
self.broadcast_stream.wait_stream(self.main_stream)
|
||||
torch.distributed.broadcast(
|
||||
drafts, src=self.last_rank, group=self.broadcast_group
|
||||
)
|
||||
drafts.record_stream(self.broadcast_stream)
|
||||
|
||||
@@ -256,7 +256,6 @@ class DFlashSpeculator(DraftModelSpeculator):
|
||||
num_tokens_across_dp=num_tokens_across_dp,
|
||||
slot_mapping=slot_mappings,
|
||||
batch_descriptor=batch_descriptor,
|
||||
is_padding=self.input_buffers.is_padding[:num_tokens],
|
||||
):
|
||||
last_hidden_states = self.model(
|
||||
input_ids=self.input_buffers.input_ids[:num_tokens],
|
||||
@@ -488,7 +487,6 @@ def _prepare_dflash_inputs_kernel(
|
||||
# Outputs
|
||||
out_input_ids_ptr,
|
||||
out_query_positions_ptr,
|
||||
out_is_padding_ptr,
|
||||
out_query_start_loc_ptr,
|
||||
out_seq_lens_ptr,
|
||||
out_query_slot_mapping_ptr,
|
||||
@@ -619,7 +617,6 @@ def _prepare_dflash_inputs_kernel(
|
||||
)
|
||||
|
||||
tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query)
|
||||
tl.store(out_is_padding_ptr + query_idx, False, mask=is_query)
|
||||
clamped_query_pos = tl.minimum(query_pos, max_model_len - 1)
|
||||
tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query)
|
||||
tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query)
|
||||
@@ -682,9 +679,6 @@ def _prepare_dflash_inputs_kernel(
|
||||
for i in range(q_pad_start, max_num_tokens, BLOCK_SIZE):
|
||||
block = i + tl.arange(0, BLOCK_SIZE)
|
||||
mask = block < max_num_tokens
|
||||
tl.store(out_input_ids_ptr + block, 0, mask=mask)
|
||||
tl.store(out_query_positions_ptr + block, 0, mask=mask)
|
||||
tl.store(out_is_padding_ptr + block, True, mask=mask)
|
||||
tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask)
|
||||
|
||||
|
||||
@@ -736,7 +730,6 @@ def prepare_dflash_inputs(
|
||||
_prepare_dflash_inputs_kernel[(num_reqs, num_blocks)](
|
||||
input_buffers.input_ids,
|
||||
input_buffers.positions,
|
||||
input_buffers.is_padding,
|
||||
input_buffers.query_start_loc,
|
||||
input_buffers.seq_lens,
|
||||
query_slot_mapping,
|
||||
|
||||
@@ -445,11 +445,6 @@ def _warmup_kernels(
|
||||
for step_indices, step_spec_flags in decode_steps:
|
||||
_run_decode_step(step_indices, step_spec_flags)
|
||||
|
||||
# The deferred PP post-update path only runs on real steps, so the steps
|
||||
# above never JIT-compile its kernel on non-last ranks.
|
||||
if not model_runner.is_last_pp_rank and model_runner.pp_handler is not None:
|
||||
model_runner.warmup_pp_decode_update()
|
||||
|
||||
# Clean up - process finish_req_ids.
|
||||
cleanup_output = SchedulerOutput.make_empty()
|
||||
cleanup_output.finished_req_ids = set(req_ids)
|
||||
|
||||
@@ -804,15 +804,6 @@ class Worker(WorkerBase):
|
||||
|
||||
@instrument(span_name="Warmup (GPU)")
|
||||
def compile_or_warm_up_model(self) -> CompilationTimes:
|
||||
# All warmup phases below run synthetic steps whose sampled outputs are
|
||||
# discarded. The PP sampled-token broadcast would carry no payload, and
|
||||
# its side-stream NCCL ops can overlap the next step's activation p2p
|
||||
# and deadlock the pipeline, so keep it disabled for the whole warmup
|
||||
# window and restore it before serving.
|
||||
pp_handler = getattr(self.model_runner, "pp_handler", None)
|
||||
if pp_handler is not None:
|
||||
pp_handler.set_disabled(True)
|
||||
|
||||
warmup_sizes: list[int] = []
|
||||
|
||||
if (
|
||||
@@ -991,9 +982,6 @@ class Worker(WorkerBase):
|
||||
# gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
|
||||
enable_gpu_sync_check()
|
||||
|
||||
if pp_handler is not None:
|
||||
pp_handler.set_disabled(False)
|
||||
|
||||
return CompilationTimes(
|
||||
language_model=self.compilation_config.compilation_time,
|
||||
encoder=self.compilation_config.encoder_compilation_time,
|
||||
|
||||
Reference in New Issue
Block a user