From e8c4c900dc3bffc5981daa560fc75d2d1985c1b7 Mon Sep 17 00:00:00 2001 From: Pi Date: Thu, 4 Jun 2026 12:15:22 -0700 Subject: [PATCH 01/37] Add moe e2e test (#2810) Signed-off-by: --- .buildkite/rl/moe_expert_ids.yml | 2 +- tests/runner/test_moe_expert_ids.py | 146 ++++++++++++++++++++++++++++ 2 files changed, 147 insertions(+), 1 deletion(-) create mode 100644 tests/runner/test_moe_expert_ids.py diff --git a/.buildkite/rl/moe_expert_ids.yml b/.buildkite/rl/moe_expert_ids.yml index 435a20db11..b9b1ae9c1c 100644 --- a/.buildkite/rl/moe_expert_ids.yml +++ b/.buildkite/rl/moe_expert_ids.yml @@ -22,7 +22,7 @@ steps: agents: queue: "${TPU_QUEUE_SINGLE:-tpu_v6e_queue}" commands: - - .buildkite/scripts/run_in_docker.sh python3 -m pytest -s -v /workspace/tpu_inference/tests/core/test_moe_expert_ids.py + - .buildkite/scripts/run_in_docker.sh python3 -m pytest -s -v /workspace/tpu_inference/tests/runner/test_moe_expert_ids.py - label: "${TPU_VERSION:-tpu6e} Record correctness test result for moe_expert_ids" key: "${TPU_VERSION:-tpu6e}_record_moe_expert_ids_CorrectnessTest" depends_on: "${TPU_VERSION:-tpu6e}_moe_expert_ids_CorrectnessTest" diff --git a/tests/runner/test_moe_expert_ids.py b/tests/runner/test_moe_expert_ids.py new file mode 100644 index 0000000000..d805d46009 --- /dev/null +++ b/tests/runner/test_moe_expert_ids.py @@ -0,0 +1,146 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import pytest +from vllm import LLM, SamplingParams + +LLM.__repr__ = lambda self: "LLM" + + +@pytest.fixture(scope="function") +def llm_enabled(): + engine = LLM( + model="Qwen/Qwen1.5-MoE-A2.7B", + load_format="dummy", + trust_remote_code=True, + max_model_len=128, + max_num_batched_tokens=128, + max_num_seqs=16, + gpu_memory_utilization=0.95, + tensor_parallel_size=1, + pipeline_parallel_size=1, + enable_prefix_caching=False, + kv_cache_dtype="auto", + enable_expert_parallel=False, + enable_return_routed_experts=True, + ) + yield engine + del engine + import gc + gc.collect() + import time + time.sleep(10) + + +@pytest.fixture(scope="function") +def llm_disabled(): + engine = LLM( + model="Qwen/Qwen1.5-MoE-A2.7B", + load_format="dummy", + trust_remote_code=True, + max_model_len=128, + max_num_batched_tokens=128, + max_num_seqs=16, + gpu_memory_utilization=0.95, + tensor_parallel_size=1, + pipeline_parallel_size=1, + enable_prefix_caching=False, + kv_cache_dtype="auto", + enable_expert_parallel=False, + enable_return_routed_experts=False, + ) + yield engine + del engine + import gc + gc.collect() + import time + time.sleep(10) + + +class TestMoEExpertIds: + """Verify that MoE routed experts are successfully returned when enabled, + and not returned when disabled. + """ + + def test_moe_expert_ids_returned_when_enabled(self, llm_enabled: LLM): + prompt = "The capital of France is" + sampling_params = SamplingParams(temperature=0, max_tokens=10) + outputs = llm_enabled.generate([prompt], sampling_params) + output = outputs[0].outputs[0] + + # Verify that routed_experts is populated and has correct shape + assert output.routed_experts is not None, ( + "MoE models must populate routed_experts when enabled") + assert len(output.routed_experts.shape) == 3, ( + f"Expected 3D expert shape, got {output.routed_experts.shape}") + + # Verify that the token dimension has size P + G - 1 + P = len(outputs[0].prompt_token_ids) + G = len(output.token_ids) + expected_len = P + G - 1 + actual_len = output.routed_experts.shape[0] + assert actual_len == expected_len, ( + f"Expected expert 0-th dim to be P + G - 1 ({expected_len}), " + f"got {actual_len}") + + def test_moe_expert_ids_not_returned_when_disabled(self, + llm_disabled: LLM): + prompt = "The capital of France is" + sampling_params = SamplingParams(temperature=0, max_tokens=10) + outputs = llm_disabled.generate([prompt], sampling_params) + output = outputs[0].outputs[0] + + # Verify that routed_experts is None + assert output.routed_experts is None, ( + "MoE models must not populate routed_experts when disabled") + + def test_moe_expert_ids_batch_isolation(self, llm_enabled: LLM): + """Verifies that the returned expert IDs for batched requests + match their reference runs when executed in isolation. + """ + prompt_a = "The capital of France is" + prompt_b = "Explain quantum computing in simple terms:" + sampling_params = SamplingParams(temperature=0, max_tokens=10) + + # 1. Run Request A isolated + outputs_a = llm_enabled.generate([prompt_a], sampling_params) + routed_experts_a = outputs_a[0].outputs[0].routed_experts + assert routed_experts_a is not None + + # 2. Run Request B isolated + outputs_b = llm_enabled.generate([prompt_b], sampling_params) + routed_experts_b = outputs_b[0].outputs[0].routed_experts + assert routed_experts_b is not None + + # 3. Run Request A & B together in a batch + outputs_batched = llm_enabled.generate([prompt_a, prompt_b], + sampling_params) + + # Verify batched outputs match their isolated runs + routed_experts_batched_a = outputs_batched[0].outputs[0].routed_experts + routed_experts_batched_b = outputs_batched[1].outputs[0].routed_experts + + assert routed_experts_batched_a is not None + assert routed_experts_batched_b is not None + + # Verify shapes are matching + assert routed_experts_batched_a.shape == routed_experts_a.shape + assert routed_experts_batched_b.shape == routed_experts_b.shape + + # Verify values are identical (batch isolation test) + np.testing.assert_array_equal(routed_experts_batched_a, + routed_experts_a) + np.testing.assert_array_equal(routed_experts_batched_b, + routed_experts_b) From 7dd4a93710362764053f8078683d81402da8c9e4 Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Thu, 4 Jun 2026 19:22:27 +0000 Subject: [PATCH 02/37] [skip ci] Update vLLM LKG to 8d9536a775a87eae5e0a1904603aa4725562765e Signed-off-by: Buildkite Bot --- .buildkite/vllm_lkg.version | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/vllm_lkg.version b/.buildkite/vllm_lkg.version index 8203508ac6..9e108debb4 100644 --- a/.buildkite/vllm_lkg.version +++ b/.buildkite/vllm_lkg.version @@ -1 +1 @@ -99ef652907932fe6ba1fb28fa8bfbafeea8ed317 +8d9536a775a87eae5e0a1904603aa4725562765e From f0c419f1bed203f319a644c5a9f379a8f88e4f86 Mon Sep 17 00:00:00 2001 From: Johnny Yang <24908445+jcyang43@users.noreply.github.com> Date: Thu, 4 Jun 2026 13:42:59 -0700 Subject: [PATCH 03/37] Update owner (#2816) Signed-off-by: Johnny Yang --- .github/CODEOWNERS | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 1b321e790c..0e566625a8 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -10,8 +10,8 @@ * @vipannalla @kyuyeunk @QiliangCui # CI/CD and Build Configuration -/.buildkite/ @jcyang43 @QiliangCui @yiw-wang @CienetStingLin -/.github/ @jcyang43 @QiliangCui @yiw-wang @CienetStingLin +/.buildkite/ @QiliangCui @yiw-wang @CienetStingLin +/.github/ @QiliangCui @yiw-wang @CienetStingLin # Documentation /docs/ @bvrockwell @vipannalla From 7c79ea02525c770ac232b7d597a6099249cd98d2 Mon Sep 17 00:00:00 2001 From: "Yuhang Ning (Daniel)" <114621152+yuhangning815@users.noreply.github.com> Date: Thu, 4 Jun 2026 13:57:20 -0700 Subject: [PATCH 04/37] GDN recurrent scan kernel: added async dma for prefill, decode, and transition (#2650) Signed-off-by: Daniel Ning Co-authored-by: Daniel Ning --- .../kernels/gdn/v2/recurrent_scan_v2.py | 484 ++++++++++++------ 1 file changed, 327 insertions(+), 157 deletions(-) diff --git a/tpu_inference/kernels/gdn/v2/recurrent_scan_v2.py b/tpu_inference/kernels/gdn/v2/recurrent_scan_v2.py index 16abac8d73..a4f84e98a7 100644 --- a/tpu_inference/kernels/gdn/v2/recurrent_scan_v2.py +++ b/tpu_inference/kernels/gdn/v2/recurrent_scan_v2.py @@ -22,6 +22,34 @@ from tpu_inference.kernels.gdn.v2 import \ compute_schedule_v2 as compute_schedule_table_v2 +# def invert_triangular_matrix(A, block_size=None): +# """Inverts a unit lower triangular matrix A using Neumann doubling. + +# Algorithm: Neumann doubling. For L strictly lower triangular of size N x N, +# since L^N = 0 we have: +# (I + L)^{-1} = (I - L)(I + L^2)(I + L^4) ... (I + L^(N/2)) + +# Args: +# A: Unit lower triangular matrix of shape (B, N, N). +# block_size: Size of the blocks for Gaussian elimination (unused). + +# Returns: +# Inverse of A, of shape (B, N, N). + +# For N=128 this is exactly 6 iterations = 12 matmuls, all (B, 128, 128). +# """ +# B, N, _ = A.shape +# num_iters = max(1, (N - 1).bit_length() - 1) +# in_dtype = A.dtype +# A_f32 = A.astype(jnp.float32) +# L = jnp.tril(A_f32, k=-1) +# eye = jnp.broadcast_to(jnp.eye(N, dtype=jnp.float32), (B, N, N)) +# Y = eye - L +# for _ in range(num_iters): +# L = jnp.matmul(L, L, precision=jax.lax.Precision.HIGHEST) +# Y = Y + jnp.matmul(Y, L, precision=jax.lax.Precision.HIGHEST) +# return Y.astype(in_dtype) + def invert_triangular_matrix(A, block_size=16): """Inverts a unit lower triangular matrix A block-wise. @@ -121,17 +149,27 @@ def inner_kernel( # VMEM scratchpad: (2, n_v, d_k, d_v). To carry state across chunks # (double buffered) prefill_scratch, - # VMEM scratchpad: (1, 2, n_v, d_k, d_v). TODO: double - # buffer or x buffer to to loop over BT in decode without overwriting state and using async copy for state load/store) + # VMEM scratchpad: (1, n_v, d_k, d_v). Per-iter safe-copy of the loaded + # state. Required as a separate buffer from decode_load_scratch because + # the prefetch DMA for iter b+2 writes to the same slot concurrently with + # this iter's compute; this buffer isolates the reads. Stored as bf16; + # per-head fp32 cast happens in VREG inside the compute loop. decode_state_scratch, - # VMEM scratchpad: (1, n_v, d_k, d_v). dtype = recurrent_state dtype - # TODO: if output dtype of state is always f32 then this can be removed. + # VMEM scratchpad: (1, n_v, d_k, d_v). Aliased to slot 0 of + # decode_store_scratch in _run_with_scratch (decode drains its stores + # before prefill runs, so the slot is free for prefill bf16 staging). state_commit_scratch, + # VMEM scratchpad: (2, n_v, d_k, d_v). Double-buffered staging for + # fully-async decode loads. iter b lands in slot (b % 2). + decode_load_scratch, + # VMEM scratchpad: (2, n_v, d_k, d_v). Double-buffered staging for + # fully-async decode stores. iter b uses slot (b % 2). + decode_store_scratch, # VMEM scratchpad: (BT, n_v * d_v). To hold decode outputs before DMA decode_output_scratch, # Array of C semaphores for decode state loads decode_read_semaphores, - # 1 semaphore for decode state stores + # 2 semaphores (one per decode_store_scratch slot) for async decode stores decode_write_semaphore, # 1 semaphore for prefill DMA (stores only) prefill_semaphore, @@ -167,7 +205,6 @@ def l2_normalize(x, eps=1e-6): return x / norm # 2. Decode Branch - # check current iteration had decode work @pl.when(decode_valid > 0) def decode_wrapper(): @@ -176,24 +213,71 @@ def get_target_idx(b): state_indices.shape[0] - 1) return state_indices[safe_req_id][...] - def process_decode(b, _): - # token by token check if decode token or not + # Pre-loop: kick off async loads for iters 0 and 1. + # iter b consumes the load that lands in decode_load_scratch[b % 2]. + # Subsequent loads (iter b+2 for each iter b) are issued from inside + # the loop as prefetches. + @pl.when(decode_count >= 1) + def _preload_slot_0(): + tgt = get_target_idx(0) + op = pltpu.make_async_copy( + src_ref=recurrent_state_in.at[pl.ds(tgt, 1)], + dst_ref=decode_load_scratch.at[pl.ds(0, 1)], + sem=decode_read_semaphores.at[0], + ) + op.start() + + @pl.when(decode_count >= 2) + def _preload_slot_1(): + tgt = get_target_idx(1) + op = pltpu.make_async_copy( + src_ref=recurrent_state_in.at[pl.ds(tgt, 1)], + dst_ref=decode_load_scratch.at[pl.ds(1, 1)], + sem=decode_read_semaphores.at[1], + ) + op.start() + + def process_decode(b, store_inflight): + # store_inflight: tuple (s0_inflight, s1_inflight) of int32 scalars. + # s{n}_inflight == 1 iff slot n has an in-flight async store DMA. + s0_inflight, s1_inflight = store_inflight is_valid = b < decode_count + slot = b % 2 + using_slot_0 = slot == 0 + cur_slot_inflight = jax.lax.select(using_slot_0, s0_inflight, + s1_inflight) @pl.when(is_valid) def do_work(): - target_idx = get_target_idx(b) - - # Load state TODO: make async - copy_op = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(target_idx, 1)], - dst_ref=state_commit_scratch, - sem=decode_read_semaphores.at[0], + # Wait for THIS iter's load (issued by preload or by iter b-2 prefetch). + wait_load = pltpu.make_async_copy( + src_ref=recurrent_state_in.at[pl.ds(0, 1)], + dst_ref=decode_load_scratch.at[pl.ds(slot, 1)], + sem=decode_read_semaphores.at[slot], ) - copy_op.start() - copy_op.wait() - decode_state_scratch[pl.ds( - 0, 1)] = state_commit_scratch[...].astype(jnp.float32) + wait_load.wait() + + # Safe-copy of loaded state. Isolates compute from the + # prefetch DMA below that writes the same slot of + # decode_load_scratch concurrently. bf16 -> bf16, no cast. + decode_state_scratch[pl.ds(0, 1)] = decode_load_scratch[pl.ds( + slot, 1)][...] + + # Prefetch load for iter b+2 (same slot, since (b+2) % 2 == b % 2). + # This DMA overlaps with the compute below. + next_b = b + 2 + + @pl.when(next_b < decode_count) + def _prefetch_next_load(): + next_tgt = get_target_idx(next_b) + op = pltpu.make_async_copy( + src_ref=recurrent_state_in.at[pl.ds(next_tgt, 1)], + dst_ref=decode_load_scratch.at[pl.ds(slot, 1)], + sem=decode_read_semaphores.at[slot], + ) + op.start() + + target_idx = get_target_idx(b) key_dim = n_kq * d_k b_aligned = (b // sublanesize) * sublanesize @@ -263,7 +347,8 @@ def do_work(): k_h = k[h:h + 1, :] # (1, d_k) v_h = v[h:h + 1, :] # (1, d_v) - state_h = current_state[h] # (d_k, d_v) + state_h = current_state[h].astype( + jnp.float32) # (d_k, d_v) k_state_h = pl.dot( k_h, state_h, @@ -312,23 +397,6 @@ def do_work(): new_state = jnp.stack(new_state_list, axis=0) # (n_v, d_k, d_v) - # TODO: remove VPU path if MXU is certified path - # decay_exp = decay[..., None] # (n_v, 1) - - # k_state = jnp.sum(k[..., None] * current_state, axis=1) # (n_v, d_v) - # v_diff = v - decay_exp * k_state - # v_new = curr_beta[..., None] * v_diff # (n_v, d_v) - - # q_state = jnp.sum(q[..., None] * current_state, axis=1) # (n_v, d_v) - # q_k = jnp.sum(q * k, axis=-1, keepdims=True) # (n_v, 1) - - # out = decay_exp * q_state + q_k * v_new # (n_v, d_v) - # k_v_new = k[..., None] * v_new[:, None, :] - # new_state = current_state * decay_exp[..., None] + k_v_new - - decode_state_scratch[pl.ds( - 0, 1)] = new_state[None, ...].astype(current_state.dtype) - # Accumulate output in scratchpad current_output = decode_output_scratch[...] mask = (jnp.arange(BT) == b).astype(current_output.dtype)[:, @@ -341,23 +409,68 @@ def do_work(): decode_output_scratch[...] = new_output.astype( current_output.dtype) - # Store state (Synchronous) - state_commit_scratch[0] = decode_state_scratch[0].astype( - state_commit_scratch.dtype) + # Async store. Before writing to decode_store_scratch[slot], + # wait for the previous same-slot store DMA (from iter b-2) + # so we don't clobber a buffer that's still being read. + @pl.when(cur_slot_inflight > 0) + def _wait_same_slot_store(): + temp_desc = pltpu.make_async_copy( + src_ref=decode_store_scratch.at[pl.ds(slot, 1)], + dst_ref=recurrent_state_out.at[pl.ds(0, 1)], + sem=decode_write_semaphore.at[slot], + ) + temp_desc.wait() + + decode_store_scratch[slot] = new_state.astype( + decode_store_scratch.dtype) copy_op = pltpu.make_async_copy( - src_ref=state_commit_scratch, + src_ref=decode_store_scratch.at[pl.ds(slot, 1)], dst_ref=recurrent_state_out.at[pl.ds(target_idx, 1)], - sem=decode_write_semaphore.at[0], + sem=decode_write_semaphore.at[slot], ) copy_op.start() - copy_op.wait() - - return None - - return None + # No wait — drained after fori_loop. + + # Update carry: this iter marked its slot as having an in-flight + # store iff is_valid. + next_s0_inflight = jax.lax.select( + is_valid & using_slot_0, + jnp.int32(1), + s0_inflight, + ) + next_s1_inflight = jax.lax.select( + is_valid & (~using_slot_0), + jnp.int32(1), + s1_inflight, + ) + return (next_s0_inflight, next_s1_inflight) # loop over bt, could be for loop, BT is static anyway, unroll - jax.lax.fori_loop(0, BT, process_decode, None) + final_s0_inflight, final_s1_inflight = jax.lax.fori_loop( + 0, + BT, + process_decode, + (jnp.int32(0), jnp.int32(0)), + ) + + # Drain any remaining async store DMAs (at most one per slot). + @pl.when(final_s0_inflight > 0) + def _drain_slot_0(): + temp_desc = pltpu.make_async_copy( + src_ref=decode_store_scratch.at[pl.ds(0, 1)], + dst_ref=recurrent_state_out.at[pl.ds(0, 1)], + sem=decode_write_semaphore.at[0], + ) + temp_desc.wait() + + @pl.when(final_s1_inflight > 0) + def _drain_slot_1(): + temp_desc = pltpu.make_async_copy( + src_ref=decode_store_scratch.at[pl.ds(1, 1)], + dst_ref=recurrent_state_out.at[pl.ds(0, 1)], + sem=decode_write_semaphore.at[1], + ) + temp_desc.wait() # Mask and write accumulated outputs to HBM mask = (jnp.arange(BT) @@ -378,29 +491,24 @@ def process_prefill(): prefill_slot = prefill_req_id % 2 def process_regular_prefill(): - # 1. Initialize state if first chunk of the request in this step - @pl.when(is_first_chunk > 0) - def init_state(): - has_init = has_initial_state[prefill_req_id][...] - - def load_from_hbm(): - state_idx = state_indices[prefill_req_id][...] - copy_op = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(state_idx, 1)], - dst_ref=state_commit_scratch, - sem=prefill_semaphore.at[prefill_slot], - ) - copy_op.start() - copy_op.wait() - prefill_scratch[prefill_slot] = state_commit_scratch[ - 0].astype(prefill_scratch.dtype) + init_has_init = has_initial_state[prefill_req_id][...] + init_state_idx = state_indices[prefill_req_id][...] + should_load_init = (is_first_chunk > 0) & (init_has_init > 0) + should_zero_init = (is_first_chunk > 0) & (init_has_init == 0) - def zero_state(): - prefill_scratch[prefill_slot] = jnp.zeros( - (n_v, d_k, d_v), dtype=prefill_scratch.dtype) + @pl.when(should_load_init) + def _start_init_load(): + copy_op = pltpu.make_async_copy( + src_ref=recurrent_state_in.at[pl.ds(init_state_idx, 1)], + dst_ref=state_commit_scratch, + sem=prefill_semaphore.at[prefill_slot], + ) + copy_op.start() - jax.lax.cond(has_init > 0, load_from_hbm, zero_state) - return None + @pl.when(should_zero_init) + def _zero_init_state(): + prefill_scratch[prefill_slot] = jnp.zeros( + (n_v, d_k, d_v), dtype=prefill_scratch.dtype) ### Preparataion for chunk wise math, ### this kernel design could be optimized lot by not doing this every chunk @@ -451,17 +559,18 @@ def zero_state(): q = l2_normalize(q) k = l2_normalize(k) - repeat_factor = n_v // n_kq - if repeat_factor > 1: - q = jnp.repeat(q, repeat_factor, axis=1) - k = jnp.repeat(k, repeat_factor, axis=1) - - # TODO: eliminate these transposes by directly slicing in the right - # shape above, + # Transpose first, then repeat: the transpose operates on the + # smaller (C, n_kq, d_k) tensor; the subsequent repeat on the + # (now-leading) axis produces the same final layout. q = q.transpose(1, 0, 2) k = k.transpose(1, 0, 2) v = v.transpose(1, 0, 2) + repeat_factor = n_v // n_kq + if repeat_factor > 1: + q = jnp.repeat(q, repeat_factor, axis=0) + k = jnp.repeat(k, repeat_factor, axis=0) + scale = d_k**-0.5 q = q * scale @@ -474,11 +583,16 @@ def zero_state(): g_cumsum = jnp.stack(g_cumsum_list, axis=-1) k_beta = k * beta[..., None] - S = jnp.matmul( - k_beta.astype(jnp.float32), - k.transpose(0, 2, 1).astype(jnp.float32), + # Fuse S and S_q into a single matmul. + k_T = k.transpose(0, 2, 1) + kbeta_q = jnp.concatenate([k_beta, q], axis=1) # (n_v, 2C, d_k) + S_both = jnp.matmul( + kbeta_q.astype(jnp.float32), + k_T.astype(jnp.float32), precision=jax.lax.Precision.HIGHEST, - ) + ) # (n_v, 2C, C) + S = S_both[:, :C, :] + S_q = S_both[:, C:, :] g_diff = g_cumsum[..., :, None] - g_cumsum[..., None, :] i = jnp.arange(C)[:, None] @@ -494,11 +608,6 @@ def zero_state(): S = jnp.where(mask_float[None, :, :] > 0, S * jnp.exp(g_diff_safe), 0.0) - S_q = jnp.matmul( - q.astype(jnp.float32), - k.transpose(0, 2, 1).astype(jnp.float32), - precision=jax.lax.Precision.HIGHEST, - ) mask_float_q = (i >= j).astype(jnp.float32) g_diff_Sq = g_diff_safe * mask_float_q[None, ...] + ( 1.0 - mask_float_q[None, ...]) * (-1e30) @@ -509,31 +618,46 @@ def zero_state(): # TODO: call the function in kernels file A_inv = invert_triangular_matrix(I_plus_S, block_size=16) - # UW + # Fuse u and w into a single matmul. Both compute + # A_inv @ ; stack v_beta and k_beta_g along the last + # axis (d_v -> d_v + d_k), do one matmul, then split. v_beta = v * beta[..., None] - u = jnp.matmul(A_inv, - v_beta.astype(jnp.float32), - precision=jax.lax.Precision.HIGHEST) - k_beta_g = k_beta * jnp.exp(g_cumsum)[..., None] - w = jnp.matmul( - A_inv, - k_beta_g.astype(jnp.float32), - precision=jax.lax.Precision.HIGHEST, - ) + vk_in = jnp.concatenate( + [ + v_beta.astype(jnp.float32), + k_beta_g.astype(jnp.float32), + ], + axis=2, + ) # (n_v, C, d_v + d_k) + uw = jnp.matmul(A_inv, vk_in, precision=jax.lax.Precision.HIGHEST) + u = uw[..., :d_v] + w = uw[..., d_v:] q_g = q * jnp.exp(g_cumsum)[..., None] + + # Fuse attn_inter and v_prime into a single matmul. With attentionDP, the async copy leads to very small perf gain. + @pl.when(should_load_init) + def _finish_init_load(): + temp_desc = pltpu.make_async_copy( + src_ref=recurrent_state_in.at[pl.ds(init_state_idx, 1)], + dst_ref=state_commit_scratch, + sem=prefill_semaphore.at[prefill_slot], + ) + temp_desc.wait() + prefill_scratch[prefill_slot] = state_commit_scratch[0].astype( + prefill_scratch.dtype) + current_state = prefill_scratch[prefill_slot] - attn_inter = jnp.matmul( - q_g.astype(jnp.float32), + qw = jnp.concatenate([q_g.astype(jnp.float32), w], axis=1) + comb = jnp.matmul( + qw, current_state.astype(jnp.float32), precision=jax.lax.Precision.HIGHEST, - ) - v_prime = jnp.matmul( - w, - current_state.astype(jnp.float32), - precision=jax.lax.Precision.HIGHEST, - ) + ) # (n_v, 2C, d_v) + attn_inter = comb[:, :C, :] + v_prime = comb[:, C:, :] + v_new = u - v_prime term2 = jnp.matmul(S_q, v_new, precision=jax.lax.Precision.HIGHEST) o_c = attn_inter + term2 @@ -552,22 +676,22 @@ def zero_state(): prefill_scratch[prefill_slot] = h_new.astype(prefill_scratch.dtype) - # Store state only if it's the last chunk of the request + # Store state only if it's the last chunk of the request. + store_state_idx = state_indices[prefill_req_id][...] + @pl.when(is_last_chunk > 0) def store_state(): # TODO: if dtype of state in HBM is always f32, # then we can eliminate this copy and directly write from scratch to HBM state_commit_scratch[0] = prefill_scratch[prefill_slot].astype( state_commit_scratch.dtype) - state_idx = state_indices[prefill_req_id][...] copy_op = pltpu.make_async_copy( src_ref=state_commit_scratch, - dst_ref=recurrent_state_out.at[pl.ds(state_idx, 1)], + dst_ref=recurrent_state_out.at[pl.ds(store_state_idx, 1)], sem=prefill_semaphore.at[prefill_slot], ) copy_op.start() copy_op.wait() - return None # TODO: eliminate this transpose and reshape by directly writing in the right shape above o_c_tr = o_c.transpose(1, 0, 2) @@ -578,6 +702,7 @@ def store_state(): o_c_flat_masked = o_c_flat * mask_float[:, None] prefill_output_ref[...] = o_c_flat_masked.astype( prefill_output_ref.dtype) + return None def process_transition_prefill(): @@ -585,6 +710,23 @@ def process_transition_prefill(): C_trans = sublanesize key_dim = n_kq * d_k + first_req_id = schedule_table[step, 11][...] + first_is_first = schedule_table[step, 11 + C_trans][...] + first_slot = first_req_id % 2 + first_has_init = has_initial_state[first_req_id][...] + should_load_first = (first_is_first > 0) & (first_has_init > 0) + first_state_idx = state_indices[first_req_id][...] + + # Async initial load + @pl.when(should_load_first) + def _start_first_load(): + copy_op = pltpu.make_async_copy( + src_ref=recurrent_state_in.at[pl.ds(first_state_idx, 1)], + dst_ref=state_commit_scratch, + sem=prefill_semaphore.at[first_slot], + ) + copy_op.start() + # Workaround: Upcast to fp32 to avoid NaNs qkv_chunk = prefill_qkv_ref[:C_trans, :].astype(jnp.float32) # Fused SiLU TODO: maybe 'SiLU' needs to be parametrized, @@ -616,50 +758,46 @@ def process_transition_prefill(): q = l2_normalize(q) k = l2_normalize(k) - repeat_factor = n_v // n_kq - if repeat_factor > 1: - q = jnp.repeat(q, repeat_factor, axis=1) - k = jnp.repeat(k, repeat_factor, axis=1) - - # TODO: eliminate these transposes by directly slicing in the right shape above, + # Transpose first, then repeat q = q.transpose(1, 0, 2) k = k.transpose(1, 0, 2) v = v.transpose(1, 0, 2) + repeat_factor = n_v // n_kq + if repeat_factor > 1: + q = jnp.repeat(q, repeat_factor, axis=0) + k = jnp.repeat(k, repeat_factor, axis=0) + scale = d_k**-0.5 q = q * scale - # state indice for req - first_req_id = schedule_table[step, 11][...] - first_is_first = schedule_table[step, 11 + C_trans][...] - first_slot = first_req_id % 2 - first_has_init = has_initial_state[first_req_id][...] - - @pl.when((first_is_first > 0) & (first_has_init > 0)) - def load_first_state(): - state_idx = state_indices[first_req_id][...] - copy_op = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(state_idx, 1)], + # Cold-start: zero the slot in place when the first sequence has + # no carried-over state, then read. + @pl.when((first_is_first > 0) & (first_has_init == 0)) + def _zero_first_slot(): + prefill_scratch[first_slot] = jnp.zeros( + (n_v, d_k, d_v), dtype=prefill_scratch.dtype) + + # Finish the async initial load started above + @pl.when(should_load_first) + def _finish_first_load(): + temp_desc = pltpu.make_async_copy( + src_ref=recurrent_state_in.at[pl.ds(first_state_idx, 1)], dst_ref=state_commit_scratch, sem=prefill_semaphore.at[first_slot], ) - copy_op.start() - copy_op.wait() + temp_desc.wait() prefill_scratch[first_slot] = state_commit_scratch[0].astype( prefill_scratch.dtype) h = prefill_scratch[first_slot] - h = jnp.where((first_is_first > 0) & (first_has_init == 0), - jnp.zeros_like(h), h) current_r = first_req_id sequence_valid = True - # loop over token by token + # Loop over tokens in the sublane. for i in range(sublanesize): - # read transition token metadata t_req = schedule_table[step, 11 + i][...] - # get sequence index for token i in sublane t_is_first = schedule_table[step, 11 + C_trans + i][...] t_is_last = schedule_table[step, 11 + 2 * C_trans + i][...] @@ -674,18 +812,15 @@ def load_first_state(): c_slot = current_r % 2 - h0 = prefill_scratch[0] - h1 = prefill_scratch[1] - prefill_scratch[0] = jnp.where(c_slot == 0, h, h0) - prefill_scratch[1] = jnp.where(c_slot == 1, h, h1) - - # prefill_scratch in f32, state_commit might be in bf16 - state_commit_scratch[0] = prefill_scratch[c_slot].astype( - state_commit_scratch.dtype) + # Commit the previous iter's h to the current request's slot. + # The other slot is untouched. + prefill_scratch[c_slot] = h def do_write(): - # TODO: Make async state_idx = state_indices[current_r][...] + # Stage h only when we actually DMA out. + state_commit_scratch[0] = h.astype( + state_commit_scratch.dtype) copy_op = pltpu.make_async_copy( src_ref=state_commit_scratch, dst_ref=recurrent_state_out.at[pl.ds(state_idx, 1)], @@ -717,13 +852,14 @@ def load_t_state(): should_load_t = (t_is_first > 0) & (t_has_init > 0) jax.lax.cond(should_load_t, load_t_state, lambda: None) - h0_new = prefill_scratch[0] - h1_new = prefill_scratch[1] - new_h = jnp.where(t_slot == 0, h0_new, h1_new) + # Cold-start: zero the slot in place for a new sequence with + # no carried-over state. Mutually exclusive with load_t_state. + @pl.when((t_is_first > 0) & (t_has_init == 0)) + def _zero_t_slot(): + prefill_scratch[t_slot] = jnp.zeros( + (n_v, d_k, d_v), dtype=prefill_scratch.dtype) - new_h = jnp.where((t_is_first > 0) & (t_has_init == 0), - jnp.zeros_like(new_h), new_h) - h = new_h + h = prefill_scratch[t_slot] current_r = t_req @@ -749,7 +885,7 @@ def load_t_state(): h = jnp.where(sequence_valid, h_new, h) - # Mask output BEFORE invalidating the sequence for the next token + # Mask output before invalidating the sequence for the next token. out_i = jnp.where(sequence_valid, out_i, 0.0) sequence_valid = jnp.where(t_is_last > 0, False, @@ -760,15 +896,16 @@ def load_t_state(): final_slot = current_r % 2 prefill_scratch[final_slot] = h - state_commit_scratch[0] = h.astype(state_commit_scratch.dtype) is_current_r_prefill = current_r >= decode_tokens - # Store state if the current request is a prefill + # Store state if the current request is a prefill. + # At the end of the transition prefill. No need to make this async. @pl.when(is_current_r_prefill) def do_final_write(): - # TODO: make async state_idx = state_indices[current_r][...] + # Stage h into state_commit only when we actually DMA out. + state_commit_scratch[0] = h.astype(state_commit_scratch.dtype) copy_op = pltpu.make_async_copy( src_ref=state_commit_scratch, dst_ref=recurrent_state_out.at[pl.ds(state_idx, 1)], @@ -978,12 +1115,17 @@ def fused_kernel( def _run_with_scratch( scratch_ref, decode_state_scratch_ref, - state_commit_scratch_ref, + decode_load_scratch_ref, + decode_store_scratch_ref, decode_output_scratch_ref, decode_read_sems, decode_write_sem, prefill_sem, ): + # Alias state_commit_scratch to slot 0 of decode_store_scratch. + # Decode drains its store DMAs before prefill runs in any step, so + # the slot is free to use as prefill's bf16 HBM staging. + state_commit_scratch_ref = decode_store_scratch_ref.at[pl.ds(0, 1)] pipeline_func = pltpu.emit_pipeline( body=functools.partial( @@ -1000,6 +1142,8 @@ def _run_with_scratch( decode_state_scratch=decode_state_scratch_ref, decode_output_scratch=decode_output_scratch_ref, state_commit_scratch=state_commit_scratch_ref, + decode_load_scratch=decode_load_scratch_ref, + decode_store_scratch=decode_store_scratch_ref, decode_read_semaphores=decode_read_sems, decode_write_semaphore=decode_write_sem, prefill_semaphore=prefill_sem, @@ -1024,7 +1168,9 @@ def _run_with_scratch( output_ref, output_ref, scratches=[ - schedule_table_ref, state_indices_ref, has_initial_state_ref + schedule_table_ref, + state_indices_ref, + has_initial_state_ref, ], ) @@ -1033,13 +1179,21 @@ def _run_with_scratch( _run_with_scratch, pltpu.VMEM((2, n_v, d_k, d_v), jnp.float32), # prefill_scratch (double buffered) - pltpu.VMEM((1, n_v, d_k, d_v), jnp.float32), # decode_state_scratch pltpu.VMEM((1, n_v, d_k, d_v), - recurrent_state_ref.dtype), # state_commit_scratch + recurrent_state_ref.dtype), # decode_state_scratch + # state_commit_scratch aliased to slot 0 of decode_store_scratch in + # _run_with_scratch; no separate allocation. + pltpu.VMEM((2, n_v, d_k, d_v), recurrent_state_ref.dtype + ), # decode_load_scratch (double-buffered) + pltpu.VMEM( + (2, n_v, d_k, d_v), recurrent_state_ref.dtype + ), # decode_store_scratch (double-buffered; slot 0 also used as prefill's state_commit staging) pltpu.VMEM((BT, n_v * d_v), mixed_qkv_ref.dtype), # decode_output_scratch - pltpu.SemaphoreType.DMA((1, )), # decode_read_semaphores - pltpu.SemaphoreType.DMA((1, )), # decode_write_semaphore + pltpu.SemaphoreType.DMA( + (2, )), # decode_read_semaphores (one per slot) + pltpu.SemaphoreType.DMA( + (2, )), # decode_write_semaphore (one per slot) pltpu.SemaphoreType.DMA((2, )), # prefill_semaphore ) @@ -1054,6 +1208,8 @@ def _run_with_scratch( "chunk_size", "BT", "use_qk_norm_in_gdn", + "vmem_limit_bytes", + "race_detect_enable", ], ) def recurrent_scan( @@ -1075,6 +1231,8 @@ def recurrent_scan( BT: int = 128, use_qk_norm_in_gdn: bool = True, has_initial_state: jax.Array | None = None, + vmem_limit_bytes: int | None = None, + race_detect_enable: bool = False, ) -> tuple[jax.Array, jax.Array]: """Fused recurrent scan kernel for GDN on TPU v7. @@ -1100,6 +1258,9 @@ def recurrent_scan( chunk_size: Block size for processing (default 128). BT: Block size for decode requests (default 128). use_qk_norm_in_gdn: Whether to use QK normalization. + vmem_limit_bytes: Per-kernel scoped VMEM ceiling passed to Mosaic. + race_detect_enable: If True, run the kernel under Pallas interpret mode with + DMA/buffer race detection enabled. Returns: A tuple containing: @@ -1115,6 +1276,10 @@ def recurrent_scan( tpu_info = pltpu.get_tpu_info() sublanesize = 4 // mixed_qkv.itemsize * tpu_info.num_sublanes + # Default the scoped VMEM ceiling. This value could be tuned for different state cache numerics and chunk sizes. + if vmem_limit_bytes is None: + vmem_limit_bytes = tpu_info.vmem_capacity_bytes + # Pad token dimension so invalid pipeline steps DMA into a safe sink area. # Sink offset must be aligned to sublanesize for Mosaic tile compatibility. block_size = max(chunk_size, BT) @@ -1185,7 +1350,12 @@ def recurrent_scan( ), grid_spec=grid_spec, input_output_aliases={1: 0}, - compiler_params=pltpu.CompilerParams(disable_bounds_checks=True), + interpret=(pltpu.InterpretParams( + detect_races=True) if race_detect_enable else False), + compiler_params=pltpu.CompilerParams( + disable_bounds_checks=True, + vmem_limit_bytes=vmem_limit_bytes, + ), )( mixed_qkv, recurrent_state, From e24c22ad88b5d2b11409d62c438de18d1770ee6d Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Thu, 4 Jun 2026 21:25:35 +0000 Subject: [PATCH 05/37] [skip ci] Update vLLM LKG to a55fccfc7cefc2d085d3557bace0e345cef67961 Signed-off-by: Buildkite Bot --- .buildkite/vllm_lkg.version | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/vllm_lkg.version b/.buildkite/vllm_lkg.version index 9e108debb4..62f663c4fd 100644 --- a/.buildkite/vllm_lkg.version +++ b/.buildkite/vllm_lkg.version @@ -1 +1 @@ -8d9536a775a87eae5e0a1904603aa4725562765e +a55fccfc7cefc2d085d3557bace0e345cef67961 From d87896c6bb8c21e16850311d2207359839a56873 Mon Sep 17 00:00:00 2001 From: Kyuyeun Kim <62023335+kyuyeunk@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:24:43 -0700 Subject: [PATCH 06/37] Update JAX to 10.1 (#2711) Signed-off-by: Kyuyeun Kim Signed-off-by: guowei-dev Co-authored-by: guowei-dev Co-authored-by: guowei-dev --- requirements.txt | 6 +- tests/e2e/test_speculative_decoding.py | 19 +++++ tests/models/jax/test_qwen3.py | 4 +- tpu_inference/env_override.py | 5 ++ .../kernels/sparse_core/ragged_gather.py | 16 ++--- .../sparse_core/ragged_gather_reduce.py | 72 ++++++++++--------- .../kernels/sparse_core/ragged_scatter.py | 56 +++++++-------- 7 files changed, 102 insertions(+), 76 deletions(-) diff --git a/requirements.txt b/requirements.txt index 194ed8f509..3d7a8a645a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,9 +5,9 @@ pytest-mock absl-py numpy google-cloud-storage -jax==0.10.0 -jaxlib==0.10.0 -libtpu==0.0.40 +jax==0.10.1 +jaxlib==0.10.1 +libtpu==0.0.41 jaxtyping flax==0.12.4 torchax==0.0.11 diff --git a/tests/e2e/test_speculative_decoding.py b/tests/e2e/test_speculative_decoding.py index 3dfeded979..1a656d705a 100644 --- a/tests/e2e/test_speculative_decoding.py +++ b/tests/e2e/test_speculative_decoding.py @@ -24,6 +24,22 @@ from vllm.v1.metrics.reader import Counter +def _disable_shardy_for_qwen35_4b(mp: pytest.MonkeyPatch) -> None: + """Disables Shardy to avoid a libtpu 0.0.41 segfault in `embed_multimodal`. + + TODO: Remove once libtpu >= 0.0.42.dev20260527. + """ + import jax + + mp.setenv("JAX_USE_SHARDY_PARTITIONER", "false") + libtpu_init_args = os.environ.get("LIBTPU_INIT_ARGS", "") + mp.setenv( + "LIBTPU_INIT_ARGS", + "--xla_use_shardy=false --xla_tpu_scoped_vmem_limit_kib=131072 " + + libtpu_init_args) + jax.config.update("jax_use_shardy_partitioner", False) + + # TODO (Qiliang Cui): remove this when XLA fixes the recursive jit call issue. def _is_v7x(): # jax.devices() will hang so use TPU_VERSION to indicate the version. @@ -520,6 +536,7 @@ def mtp_baseline(): } test_prompts = get_eagle3_test_prompts() with pytest.MonkeyPatch.context() as mp: + _disable_shardy_for_qwen35_4b(mp) ref_outputs = _get_baseline_results( mp, sampling_config, @@ -553,6 +570,7 @@ def test_mtp_correctness( model_name = "Qwen/Qwen3.5-4B" monkeypatch.setenv("MODEL_IMPL_TYPE", "vllm") monkeypatch.setenv("DRAFT_MODEL_IMPL_TYPE", "vllm") + _disable_shardy_for_qwen35_4b(monkeypatch) speculative_config = { "method": "mtp", @@ -591,6 +609,7 @@ def test_mtp_performance( model_name = "Qwen/Qwen3.5-4B" monkeypatch.setenv("MODEL_IMPL_TYPE", "vllm") monkeypatch.setenv("DRAFT_MODEL_IMPL_TYPE", "vllm") + _disable_shardy_for_qwen35_4b(monkeypatch) extra_kwargs = { "seed": 42, diff --git a/tests/models/jax/test_qwen3.py b/tests/models/jax/test_qwen3.py index ad8a020ad7..95dfa21c18 100644 --- a/tests/models/jax/test_qwen3.py +++ b/tests/models/jax/test_qwen3.py @@ -246,8 +246,8 @@ def test_expected_error_with_tight_threshold( with assert_weight_loading_memory_bounded( model, description=f"load_weights({model_name})", - threshold_multiplier=0.001, - min_threshold_bytes=1, + threshold_multiplier=-0.001, + min_threshold_bytes=-1, ), set_current_vllm_config(config): loader.load_weights(model, config.model_config) diff --git a/tpu_inference/env_override.py b/tpu_inference/env_override.py index 7066cc0e79..46c2334d54 100644 --- a/tpu_inference/env_override.py +++ b/tpu_inference/env_override.py @@ -17,6 +17,11 @@ os.environ["XLA_FLAGS"] = "--xla_cpu_max_isa=AVX2 " + os.environ.get( "XLA_FLAGS", "") +# TODO: Remove this when SMEM capacity optimization for batched rpa lands. +os.environ[ + "LIBTPU_INIT_ARGS"] = "--xla_tpu_use_dynamic_smem_negotiation=true " + os.environ.get( + "LIBTPU_INIT_ARGS", "") + # Monkeypatch vLLM to avoid ImportError: cannot import name 'SamplingParams' from 'vllm' # in vllm/v1/... submodules due to circular imports or lazy loading failures. try: diff --git a/tpu_inference/kernels/sparse_core/ragged_gather.py b/tpu_inference/kernels/sparse_core/ragged_gather.py index 22cf63b448..3482bcf9d3 100644 --- a/tpu_inference/kernels/sparse_core/ragged_gather.py +++ b/tpu_inference/kernels/sparse_core/ragged_gather.py @@ -280,19 +280,19 @@ def ragged_gather(x: jax.Array, indices: jax.Array, start: jax.Array, core_axis_name=vector_mesh.core_axis_name, subcore_axis_name=vector_mesh.subcore_axis_name, ), - out_shape=jax.ShapeDtypeStruct( + out_type=jax.ShapeDtypeStruct( (out_size + out_pad_size, aligned_hidden_size), dtype), compiler_params=pltpu.CompilerParams( use_tc_tiling_on_sc=True, disable_bounds_checks=True, ), - scratch_shapes=[ - pltpu.VMEM((num_simd_lanes, ), jnp.int32), - pltpu.VMEM((num_simd_lanes, ), jnp.int32), - pltpu.VMEM((num_simd_lanes, col_size), jnp.uint32), - pltpu.VMEM((num_simd_lanes, ), jnp.int32), - pltpu.SemaphoreType.DMA((2, )), - ], + scratch_types=dict( + start_vmem_ref=pltpu.VMEM((num_simd_lanes, ), jnp.int32), + end_vmem_ref=pltpu.VMEM((num_simd_lanes, ), jnp.int32), + out_vmem_ref=pltpu.VMEM((num_simd_lanes, col_size), jnp.uint32), + indices_vmem_ref=pltpu.VMEM((num_simd_lanes, ), jnp.int32), + sem_ref=pltpu.SemaphoreType.DMA((2, )), + ), mesh=vector_mesh, name="sc_ragged_gather", )(start, end, x, indices)[:out_size, :hidden_size] diff --git a/tpu_inference/kernels/sparse_core/ragged_gather_reduce.py b/tpu_inference/kernels/sparse_core/ragged_gather_reduce.py index db71bc8f38..0de41e5c23 100644 --- a/tpu_inference/kernels/sparse_core/ragged_gather_reduce.py +++ b/tpu_inference/kernels/sparse_core/ragged_gather_reduce.py @@ -443,30 +443,30 @@ def ragged_gather_reduce( ) -> jax.Array: """Gathers `x` according to `indices`, applies weights and masks, and reduces. - This function performs a gathered lookup from `x` using `indices`, scales the - obtained rows by `topk_weights`, masks out any rows where `valid_rows_mask` is - False, and then groups every `reduce_group_size` rows together and reduces - them via summation. - - The typical use case of this kernel is unpermute + local-reduction in the - MOE after GMM. Compared to maxtext.src.maxtext.kernels.gather_reduce_sc, - this kernel provides better performance if large sparsity exists in - `valid_rows_mask`. For example, expert_parallelism =8, 16 etc. - - Args: - x: A 2D JAX array of input features with shape `(input_size, hidden_size)`. - indices: A 1D JAX array of indices to gather with shape `(input_size,)`. - topk_weights: A 1D JAX array of weights to scale the gathered rows with - shape `(input_size,)`. - valid_rows_mask: A 1D boolean JAX array indicating which gathered rows are - valid, with shape `(input_size,)`. - reduce_group_size: An integer representing the number of consecutive rows to - reduce (sum) together. - - Returns: - A 2D JAX array of reduced data with shape - `(input_size // reduce_group_size, hidden_size)`. - """ + This function performs a gathered lookup from `x` using `indices`, scales the + obtained rows by `topk_weights`, masks out any rows where `valid_rows_mask` is + False, and then groups every `reduce_group_size` rows together and reduces + them via summation. + + The typical use case of this kernel is unpermute + local-reduction in the + MOE after GMM. Compared to maxtext.src.maxtext.kernels.gather_reduce_sc, + this kernel provides better performance if large sparsity exists in + `valid_rows_mask`. For example, expert_parallelism =8, 16 etc. + + Args: + x: A 2D JAX array of input features with shape `(input_size, hidden_size)`. + indices: A 1D JAX array of indices to gather with shape `(input_size,)`. + topk_weights: A 1D JAX array of weights to scale the gathered rows with + shape `(input_size,)`. + valid_rows_mask: A 1D boolean JAX array indicating which gathered rows are + valid, with shape `(input_size,)`. + reduce_group_size: An integer representing the number of consecutive rows to + reduce (sum) together. + + Returns: + A 2D JAX array of reduced data with shape + `(input_size // reduce_group_size, hidden_size)`. + """ assert x.ndim == 2, "ragged_gather_reduce only supports 2d inputs." assert indices.ndim == 1, "ragged_gather_reduce only supports 1d indices." @@ -553,7 +553,7 @@ def ragged_gather_reduce( num_row_partitions=num_row_partitions, num_column_partitions=num_column_partitions, ), - out_shape=jax.ShapeDtypeStruct( + out_type=jax.ShapeDtypeStruct( (x.shape[0] // reduce_group_size, x.shape[1]), jnp.float32, ), @@ -561,16 +561,18 @@ def ragged_gather_reduce( use_tc_tiling_on_sc=True, disable_bounds_checks=True, ), - scratch_shapes=[ - pltpu.VMEM((num_simd_lanes, ), jnp.int32), - pltpu.VMEM((num_simd_lanes, col_size), jnp.uint32), - pltpu.VMEM((1, col_size), jnp.uint32), - pltpu.VMEM((num_simd_lanes, ), jnp.int32), - pltpu.VMEM((num_simd_lanes, ), jnp.int32), - pltpu.VMEM((num_simd_lanes, ), jnp.float32), - pltpu.VMEM((num_simd_lanes, ), jnp.int32), - pltpu.SemaphoreType.DMA((2, )), - ], + scratch_types=dict( + num_rows_per_row_partition_vmem_ref=pltpu.VMEM((num_simd_lanes, ), + jnp.int32), + out_vmem_ref=pltpu.VMEM((num_simd_lanes, col_size), jnp.uint32), + prev_iter_last_row_vmem_ref=pltpu.VMEM((1, col_size), jnp.uint32), + src_indices_vmem_ref=pltpu.VMEM((num_simd_lanes, ), jnp.int32), + dst_indices_vmem_ref=pltpu.VMEM((num_simd_lanes, ), jnp.int32), + topk_weights_vmem_ref=pltpu.VMEM((num_simd_lanes, ), jnp.float32), + sorted_by_validity_vmem_ref=pltpu.VMEM((num_simd_lanes, ), + jnp.int32), + sem_ref=pltpu.SemaphoreType.DMA((2, )), + ), mesh=vector_mesh, name="sc_ragged_gather_reduce", )( diff --git a/tpu_inference/kernels/sparse_core/ragged_scatter.py b/tpu_inference/kernels/sparse_core/ragged_scatter.py index 59dbec2009..6ca6f1eca8 100644 --- a/tpu_inference/kernels/sparse_core/ragged_scatter.py +++ b/tpu_inference/kernels/sparse_core/ragged_scatter.py @@ -353,26 +353,26 @@ def ragged_scatter(x: jax.Array, indices: jax.Array, start: jax.Array, end: jax.Array) -> jax.Array: """Gathers rows from `x` according to `indices` within a specified range. - This function performs a gather operation equivalent to `x[indices]` for - indices that fall within the range `[start, end)`. For indices outside this - range, the behavior is undefined. - - Args: - x: A 2D JAX array to gather data from, with shape `(num_rows, hidden_size)`. - indices: A 1D JAX array of indices to gather, with shape `(output_size,)`. - start: A scalar or 1D array of size 1 containing the start index (inclusive) - to process. - end: A scalar or 1D array of size 1 containing the end index (exclusive) to - process. - - Returns: - A 2D JAX array of gathered data with shape `(output_size, hidden_size)`. - - The typical usage of this kernel is "unpermute" after GMM of the MOE layers. - That is, replace `gmm2_res[topk_argsort_revert_indices]` with - `ragged_scatter(x, topk_argsort_revert_indices, ..)` in - tpu_inference/layers/common/fused_moe_gmm.py. - """ + This function performs a gather operation equivalent to `x[indices]` for + indices that fall within the range `[start, end)`. For indices outside this + range, the behavior is undefined. + + Args: + x: A 2D JAX array to gather data from, with shape `(num_rows, hidden_size)`. + indices: A 1D JAX array of indices to gather, with shape `(output_size,)`. + start: A scalar or 1D array of size 1 containing the start index (inclusive) + to process. + end: A scalar or 1D array of size 1 containing the end index (exclusive) to + process. + + Returns: + A 2D JAX array of gathered data with shape `(output_size, hidden_size)`. + + The typical usage of this kernel is "unpermute" after GMM of the MOE layers. + That is, replace `gmm2_res[topk_argsort_revert_indices]` with + `ragged_scatter(x, topk_argsort_revert_indices, ..)` in + tpu_inference/layers/common/fused_moe_gmm.py. + """ assert x.ndim == 2, "Ragged scatter only supports 2d inputs." assert indices.ndim == 1, "Ragged scatter only supports 1d indices." @@ -432,19 +432,19 @@ def ragged_scatter(x: jax.Array, indices: jax.Array, start: jax.Array, core_axis_name=vector_mesh.core_axis_name, subcore_axis_name=vector_mesh.subcore_axis_name, ), - out_shape=jax.ShapeDtypeStruct( + out_type=jax.ShapeDtypeStruct( (out_size + out_pad_size, aligned_hidden_size), dtype), compiler_params=pltpu.CompilerParams( use_tc_tiling_on_sc=True, disable_bounds_checks=True, ), - scratch_shapes=[ - pltpu.VMEM((num_simd_lanes, ), jnp.int32), # total_num_rows - pltpu.VMEM((num_simd_lanes, col_size), jnp.uint32), - pltpu.VMEM((num_simd_lanes, ), jnp.int32), # src_indices - pltpu.VMEM((num_simd_lanes, ), jnp.int32), # dst_indices - pltpu.SemaphoreType.DMA((2, )), - ], + scratch_types=dict( + total_num_rows_vmem_ref=pltpu.VMEM((num_simd_lanes, ), jnp.int32), + out_vmem_ref=pltpu.VMEM((num_simd_lanes, col_size), jnp.uint32), + src_indices_vmem_ref=pltpu.VMEM((num_simd_lanes, ), jnp.int32), + dst_indices_vmem_ref=pltpu.VMEM((num_simd_lanes, ), jnp.int32), + sem_ref=pltpu.SemaphoreType.DMA((2, )), + ), mesh=vector_mesh, name="sc_ragged_scatter", )(total_num_rows, x, src_indices, dst_indices)[:out_size, :hidden_size] From eda5b34d786dfa3e647d340d2ba7b9353e1c4012 Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Fri, 5 Jun 2026 01:46:40 +0000 Subject: [PATCH 07/37] [skip ci] Update vLLM LKG to b7c5baf63d5f95d18493a7fae3c75fc4853732a7 Signed-off-by: Buildkite Bot --- .buildkite/vllm_lkg.version | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/vllm_lkg.version b/.buildkite/vllm_lkg.version index 62f663c4fd..e125dce9c7 100644 --- a/.buildkite/vllm_lkg.version +++ b/.buildkite/vllm_lkg.version @@ -1 +1 @@ -a55fccfc7cefc2d085d3557bace0e345cef67961 +b7c5baf63d5f95d18493a7fae3c75fc4853732a7 From cd97f44589ce7f51e972e01364a94811e2c64461 Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Fri, 5 Jun 2026 03:52:21 +0000 Subject: [PATCH 08/37] [skip ci] Update vLLM LKG to 063ce98fb7104dba93f72d56c094fb8b708bd793 Signed-off-by: Buildkite Bot --- .buildkite/vllm_lkg.version | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/vllm_lkg.version b/.buildkite/vllm_lkg.version index e125dce9c7..fd0dedc5a5 100644 --- a/.buildkite/vllm_lkg.version +++ b/.buildkite/vllm_lkg.version @@ -1 +1 @@ -b7c5baf63d5f95d18493a7fae3c75fc4853732a7 +063ce98fb7104dba93f72d56c094fb8b708bd793 From d502d4a993b385358d364df8177d0ec049cdb2ed Mon Sep 17 00:00:00 2001 From: Lumosis <30372757+Lumosis@users.noreply.github.com> Date: Thu, 4 Jun 2026 21:39:02 -0700 Subject: [PATCH 09/37] [Spec Decoding] Support Logprobs with Speculative Decoding (#2799) Signed-off-by: Lihao Ran Co-authored-by: Lihao Ran --- tests/e2e/test_speculative_decoding.py | 54 +++++ tests/spec_decode/test_utils.py | 88 ++++++- .../models/vllm/vllm_model_wrapper.py | 10 +- tpu_inference/platforms/tpu_platform.py | 15 ++ tpu_inference/runner/compilation_manager.py | 136 ++++++++++- tpu_inference/runner/tpu_runner.py | 90 ++++++- tpu_inference/runner/utils.py | 1 - tpu_inference/spec_decode/jax/utils.py | 225 ++++++++++++++++++ 8 files changed, 601 insertions(+), 18 deletions(-) diff --git a/tests/e2e/test_speculative_decoding.py b/tests/e2e/test_speculative_decoding.py index 1a656d705a..575cf8a214 100644 --- a/tests/e2e/test_speculative_decoding.py +++ b/tests/e2e/test_speculative_decoding.py @@ -202,6 +202,14 @@ def _test_correctness_helper( **kwargs) spec_outputs = spec_llm.generate(test_prompts, sampling_config) + if sampling_config.logprobs is not None: + for spec_output in spec_outputs: + completion = spec_output.outputs[0] + assert completion.logprobs is not None, "Logprobs should not be None" + assert len(completion.logprobs) == len(completion.token_ids), ( + f"Length mismatch: len(logprobs)={len(completion.logprobs)} vs " + f"len(token_ids)={len(completion.token_ids)}") + matches = 0 misses = 0 for ref_output, spec_output in zip(ref_outputs, spec_outputs): @@ -432,6 +440,7 @@ def eagle3_baseline(): "async_scheduling, enable_dp_attention", [ (False, False), + (False, True), (True, False), (True, True), ], @@ -634,3 +643,48 @@ def test_mtp_performance( model_name=model_name, extra_kwargs=extra_kwargs, ) + + +def test_eagle3_logprobs_correctness(monkeypatch: pytest.MonkeyPatch, ): + """Reproduction test for speculative decoding with logprobs enabled.""" + model_name = 'meta-llama/Meta-Llama-3-8B-Instruct' + + model_impl = os.environ.get("MODEL_IMPL_TYPE", "auto") + monkeypatch.setenv("DRAFT_MODEL_IMPL_TYPE", model_impl) + + speculative_config = { + 'model': "unkmaster/EAGLE3-LLaMA3.1-Instruct-8B", + "num_speculative_tokens": 3, + "method": "eagle3", + "draft_tensor_parallel_size": 1 + } + + sampling_config = SamplingParams(temperature=0, + max_tokens=8, + ignore_eos=True, + repetition_penalty=1, + frequency_penalty=0, + presence_penalty=0, + min_p=0, + logprobs=1) + test_prompts = [ + "Predict the continuation of this sequence: 1 2 3 4 5 6 7 8" + ] + + # Get baseline outputs with logprobs enabled + ref_outputs = _get_baseline_results(monkeypatch, + sampling_config, + model_name, + test_prompts, + max_num_seqs=2) + + # Get speculative decoding outputs with logprobs enabled + _test_correctness_helper(monkeypatch, + sampling_config, + model_name, + speculative_config, + test_prompts, + ref_outputs=ref_outputs, + max_num_seqs=2, + async_scheduling=True, + enable_dp_attention=False) diff --git a/tests/spec_decode/test_utils.py b/tests/spec_decode/test_utils.py index 8dacb89e65..54c3c4689e 100644 --- a/tests/spec_decode/test_utils.py +++ b/tests/spec_decode/test_utils.py @@ -20,7 +20,8 @@ from tpu_inference.layers.jax.sample.rejection_sampler import RejectionSampler from tpu_inference.runner.utils import SpecDecodeMetadata from tpu_inference.spec_decode.jax.utils import (PLACEHOLDER_TOKEN_ID, - extract_last_sampled_tokens) + extract_last_sampled_tokens, + filter_speculative_logprobs) VOCAB_SIZE = 100 @@ -143,3 +144,88 @@ def test_extract_last_sampled_tokens_matches_parse_output( np.testing.assert_array_equal(np.asarray(last_sampled), ref_last) np.testing.assert_array_equal(np.asarray(num_rejected), ref_num_rej) + + +def test_filter_speculative_logprobs(): + # Setup inputs matching the manual trace + dp_size = 2 + num_reqs = 3 + vocab_size = 100 + padded_tokens_length = 5 + + # log_token_ids: shape [14, 1] + log_token_ids = np.array( + [ + # Rank 0 (draft) + [10], + [11], + [12], + [20], + [-1], + # Rank 0 (bonus) + [13], + [-1], + # Rank 1 (draft) + [30], + [31], + [-1], + [-1], + [-1], + # Rank 1 (bonus) + [-1], + [-1] + ], + dtype=np.int32) + + # logprobs_arr: shape [14, 1] + logprobs_arr = np.array( + [[110.0], [111.0], [112.0], [120.0], [0.0], [113.0], [0.0], [130.0], + [131.0], [0.0], [0.0], [0.0], [0.0], [0.0]], + dtype=np.float32) + + # selected_token_ranks: shape [14] + selected_token_ranks = np.array([1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0, 0], + dtype=np.int32) + + metadata = SpecDecodeMetadata( + draft_lengths=jnp.zeros(num_reqs), # Not used by the function + target_logits_indices=jnp.zeros(1), # Not used + bonus_logits_indices=jnp.zeros(1), # Not used + final_logits_indices=jnp.zeros(padded_tokens_length * + dp_size), # Shape is used: [10] + ) + metadata.draft_lengths_cpu = np.array([3, 2, 4, 0], dtype=np.int32) + metadata.req_indices_dp = { + 0: [0, 1], + 1: [2, 3] # 3 is padding req (>= num_reqs) + } + + # Call the function + (filtered_token_ids, filtered_logprobs, filtered_ranks, + cu_num_generated_tokens) = filter_speculative_logprobs( + log_token_ids=log_token_ids, + logprobs_arr=logprobs_arr, + selected_token_ranks=selected_token_ranks, + spec_decode_metadata=metadata, + vocab_size=vocab_size, + dp_size=dp_size, + num_reqs=num_reqs, + ) + + # Expected outputs + expected_token_ids = np.array([[10], [11], [12], [13], [20], [30], [31]], + dtype=np.int32) + + expected_logprobs = np.array( + [[110.0], [111.0], [112.0], [113.0], [120.0], [130.0], [131.0]], + dtype=np.float32) + + expected_ranks = np.array([1, 1, 1, 1, 1, 1, 1], dtype=np.int32) + + expected_cu = [0, 4, 5, 7] + + # Assertions + np.testing.assert_array_equal(filtered_token_ids, expected_token_ids) + np.testing.assert_array_equal(filtered_logprobs, expected_logprobs) + np.testing.assert_array_equal(filtered_ranks, expected_ranks) + assert cu_num_generated_tokens == expected_cu diff --git a/tpu_inference/models/vllm/vllm_model_wrapper.py b/tpu_inference/models/vllm/vllm_model_wrapper.py index 6e0ea4aac3..e5c7b6bc01 100644 --- a/tpu_inference/models/vllm/vllm_model_wrapper.py +++ b/tpu_inference/models/vllm/vllm_model_wrapper.py @@ -330,10 +330,16 @@ def load_weights(self, for name, param in vllm_model.named_parameters(): full_name = f"vllm_model.{name}" if full_name in shared_params: + target_param = shared_params[full_name] + if param.shape != target_param.shape: + logger.warning( + f"Shape mismatch for shared parameter {full_name}: " + f"draft {param.shape} vs target {target_param.shape}. " + "Skipping sharing.") + continue logger.info(f"Sharing parameter: {full_name}") # torch_view creates a torchax tensor sharing memory with the JAX array - param.data = torchax.interop.torch_view( - shared_params[full_name]) + param.data = torchax.interop.torch_view(target_param) if self.vllm_config.speculative_config and self.vllm_config.speculative_config.method == "eagle3" and not self.is_draft_model: set_eagle3_aux_hidden_state_layers( diff --git a/tpu_inference/platforms/tpu_platform.py b/tpu_inference/platforms/tpu_platform.py index 8203c23cc9..22f0fe8523 100644 --- a/tpu_inference/platforms/tpu_platform.py +++ b/tpu_inference/platforms/tpu_platform.py @@ -8,6 +8,21 @@ import numpy import torch import vllm.envs as vllm_envs + +# Monkeypatch torch.accelerator.empty_cache to ignore device_allocator error on TPU. +if hasattr(torch, "accelerator") and hasattr(torch.accelerator, "empty_cache"): + _orig_empty_cache = torch.accelerator.empty_cache + + def _patched_empty_cache(*args, **kwargs): + try: + _orig_empty_cache(*args, **kwargs) + except RuntimeError as e: + if "Allocator for jax is not a DeviceAllocator" in str(e): + pass + else: + raise e + + torch.accelerator.empty_cache = _patched_empty_cache from vllm.platforms.interface import Platform, PlatformEnum from tpu_inference import envs diff --git a/tpu_inference/runner/compilation_manager.py b/tpu_inference/runner/compilation_manager.py index 08921a9c4b..164f4f9954 100644 --- a/tpu_inference/runner/compilation_manager.py +++ b/tpu_inference/runner/compilation_manager.py @@ -34,8 +34,10 @@ JaxIntermediateTensors from tpu_inference.runner.utils import SpecDecodeMetadata from tpu_inference.spec_decode.jax.utils import ( - concat_last_sampled_tokens_and_draft_tokens, extract_last_sampled_tokens) -from tpu_inference.utils import device_array, to_jax_dtype + concat_last_sampled_tokens_and_draft_tokens, extend_logits_simple, + extract_last_sampled_tokens, process_and_extend_logits) +from tpu_inference.utils import (device_array, get_mesh_shape_product, + to_jax_dtype) if TYPE_CHECKING: from tpu_inference.runner.tpu_runner import TPUModelRunner @@ -820,6 +822,33 @@ def _precompile_gather_logprobs(self) -> None: num_reqs=num_reqs, ) + if self.runner.speculative_config: + logger.info( + "Compiling gather_logprobs for speculative decoding shapes.") + for num_logits in self.runner.num_logits_paddings: + for num_reqs in self.runner.num_reqs_paddings: + if num_reqs > num_logits: + continue + combined_size = num_logits + num_reqs + logits_sharding = NamedSharding(self.runner.mesh, + PartitionSpec()) + token_ids_sharding = NamedSharding( + self.runner.mesh, + PartitionSpec(ShardingAxisName.ATTN_DATA)) + logits = self._create_dummy_tensor( + (combined_size, hsize), jnp.float32, logits_sharding) + token_ids = self._create_dummy_tensor( + (combined_size, ), jnp.int32, token_ids_sharding) + self._run_compilation( + f"worker{self.runner.rank} gather_logprobs_spec", + compute_and_gather_logprobs, + logits, + token_ids, + self.runner.model_config.max_logprobs, + num_logits=num_logits, + num_reqs=num_reqs, + ) + logger.info( "Compiling compute_and_gather_prompt_logprobs with different input shapes." ) @@ -852,12 +881,115 @@ def _precompile_gather_logprobs(self) -> None: self._gather_logprobs_precompiled = True + def _precompile_process_and_extend_logits(self) -> None: + logger.info( + "Compiling _process_and_extend_logits with different input shapes." + ) + vocab_size = self.runner.vocab_size + for num_logits in self.runner.num_logits_paddings: + for num_reqs in self.runner.num_reqs_paddings: + if num_reqs > num_logits: + continue + + logits_sharding = NamedSharding(self.runner.mesh, + PartitionSpec()) + dp_sharding = NamedSharding(self.runner.mesh, PartitionSpec()) + + target_logits = self._create_dummy_tensor( + (num_logits, vocab_size), jnp.float32, logits_sharding) + + processed_bonus_logits = self._create_dummy_tensor( + (num_reqs, vocab_size), jnp.float32, logits_sharding) + + draft_lengths = self._create_dummy_tensor( + (num_reqs, ), jnp.int32, dp_sharding) + + temperature = self._create_dummy_tensor( + (num_reqs, ), np.float32, dp_sharding) + top_k = self._create_dummy_tensor((num_reqs, ), np.int32, + dp_sharding) + top_p = self._create_dummy_tensor((num_reqs, ), np.float32, + dp_sharding) + + dummy_shape = (1, ) # logprobs=True + _cache_collision_dummy = jnp.zeros(dummy_shape, + dtype=jnp.int32) + _cache_collision_dummy = device_array(self.runner.mesh, + _cache_collision_dummy) + + sampling_metadata = TPUSupportedSamplingMetadata( + temperature=temperature, + top_k=top_k, + top_p=top_p, + _cache_collision_dummy=_cache_collision_dummy, + do_sampling=True, + logprobs=True, + ) + + spec_decode_metadata = SpecDecodeMetadata( + draft_lengths=draft_lengths, + target_logits_indices=self._create_dummy_tensor( + (num_logits, ), jnp.int32, dp_sharding), + bonus_logits_indices=self._create_dummy_tensor( + (num_reqs, ), jnp.int32, dp_sharding), + final_logits_indices=self._create_dummy_tensor( + (num_logits, ), jnp.int32, dp_sharding), + ) + + self._run_compilation( + f"worker{self.runner.rank} _process_and_extend_logits", + process_and_extend_logits, + self.runner.mesh, + target_logits, + processed_bonus_logits, + spec_decode_metadata, + sampling_metadata, + num_logits=num_logits, + num_reqs=num_reqs, + ) + + def _precompile_extend_logits_simple(self) -> None: + logger.info( + "Compiling _extend_logits_simple with different input shapes.") + vocab_size = self.runner.vocab_size + for num_logits in self.runner.num_logits_paddings: + for num_reqs in self.runner.num_reqs_paddings: + if num_reqs > num_logits: + continue + + attn_data_size = get_mesh_shape_product( + self.runner.mesh, ShardingAxisName.ATTN_DATA) + if attn_data_size == 1: + logits_spec = PartitionSpec() + else: + logits_spec = PartitionSpec(ShardingAxisName.ATTN_DATA, + None) + + logits_sharding = NamedSharding(self.runner.mesh, logits_spec) + + target_logits = self._create_dummy_tensor( + (num_logits, vocab_size), jnp.bfloat16, logits_sharding) + bonus_logits = self._create_dummy_tensor( + (num_reqs, vocab_size), jnp.bfloat16, logits_sharding) + + self._run_compilation( + f"worker{self.runner.rank} _extend_logits_simple", + extend_logits_simple, + target_logits, + bonus_logits, + self.runner.mesh, + num_logits=num_logits, + num_reqs=num_reqs, + ) + def _precompile_speculative_decoding(self) -> None: logger.info( "Compiling speculative_decoding with different input shapes.") self._precompile_rejection_sampler() self._precompile_extract_last_sampled_tokens() self._precompile_extract_draft_token_ids() + self._precompile_process_and_extend_logits() + self._precompile_extend_logits_simple() if self.runner.speculative_config.method == "eagle3": self._precompile_eagle3_helpers() if self.runner.speculative_config.method == "mtp": diff --git a/tpu_inference/runner/tpu_runner.py b/tpu_inference/runner/tpu_runner.py index dc71829d17..2bda0331da 100644 --- a/tpu_inference/runner/tpu_runner.py +++ b/tpu_inference/runner/tpu_runner.py @@ -84,7 +84,9 @@ StructuredDecodingManager from tpu_inference.spec_decode.jax.eagle3 import Eagle3Proposer from tpu_inference.spec_decode.jax.utils import ( - concat_last_sampled_tokens_and_draft_tokens, extract_last_sampled_tokens) + concat_last_sampled_tokens_and_draft_tokens, extend_logits_simple, + extract_last_sampled_tokens, filter_speculative_logprobs, + process_and_extend_logits) from tpu_inference.utils import (device_array, make_optimized_mesh, time_function, to_jax_dtype, to_torch_dtype) @@ -151,7 +153,11 @@ def get_output(self) -> ModelRunnerOutput: if self._logprobs_tensors is not None: # Use materialize to ensure logprobs are ready on host when we return async results self._model_runner_output.logprobs = _jax_logprobs_materialize( - self._logprobs_tensors, self.logits_indices_selector) + self._logprobs_tensors, + self.logits_indices_selector, + spec_decode_metadata=self._spec_decode_metadata, + runner=self._runner, + num_reqs=self._num_reqs) if self._prompt_logprobs_async_data is not None: self._model_runner_output.prompt_logprobs_dict = ( @@ -252,6 +258,7 @@ def _substitute_placeholder_token( new_token_values = next_tokens[token_in_tpu_pre_next_tokens_indices] original_values = input_ids[token_in_tpu_cur_input_indices] update_values = jnp.where(mask, new_token_values, original_values) + return input_ids.at[token_in_tpu_cur_input_indices].set(update_values) @@ -281,7 +288,10 @@ def _subtract_num_rejected_tokens_fn(seq_lens: jax.Array, positions: jax.Array, def _jax_logprobs_materialize( logprobs_tensors: LogprobsTensors, logits_indices_selector: Optional[List[int]] = None, - cu_num_generated_tokens: Optional[Any] = None) -> LogprobsLists: + cu_num_generated_tokens: Optional[Any] = None, + spec_decode_metadata: Optional[SpecDecodeMetadata] = None, + runner: Optional[Any] = None, + num_reqs: Optional[int] = None) -> LogprobsLists: """Materializes logprobs from JAX arrays into NumPy-backed LogprobsLists.""" log_token_ids = np.asarray( jax.device_get(logprobs_tensors.logprob_token_ids)) @@ -289,10 +299,43 @@ def _jax_logprobs_materialize( selected_token_ranks = np.asarray( jax.device_get(logprobs_tensors.selected_token_ranks)) - if logits_indices_selector is not None: - log_token_ids = log_token_ids[logits_indices_selector] - logprobs_arr = logprobs_arr[logits_indices_selector] - selected_token_ranks = selected_token_ranks[logits_indices_selector] + # For speculative decoding, we need to filter and reorganize the materialized + # logprobs. The raw logprobs contain info for all proposed draft tokens (including + # rejected ones) and bonus tokens. We filter them to only keep the accepted + # draft tokens and the actual bonus token for each request, and flatten them + # back to match the output format. + if (spec_decode_metadata is not None + and np.sum(spec_decode_metadata.draft_lengths_cpu) > 0): + assert runner is not None + vocab_size = runner.input_batch.vocab_size + dp_size = runner.dp_size + num_reqs = runner.input_batch.num_reqs if num_reqs is None else num_reqs + + ( + log_token_ids, + logprobs_arr, + selected_token_ranks, + cu_num_generated_tokens, + ) = filter_speculative_logprobs( + log_token_ids, + logprobs_arr, + selected_token_ranks, + spec_decode_metadata, + vocab_size, + dp_size, + num_reqs, + ) + + else: + if logits_indices_selector is not None: + log_token_ids = log_token_ids[logits_indices_selector] + logprobs_arr = logprobs_arr[logits_indices_selector] + selected_token_ranks = selected_token_ranks[ + logits_indices_selector] + + if cu_num_generated_tokens is None and runner is not None: + num_reqs = runner.input_batch.num_reqs if num_reqs is None else num_reqs + cu_num_generated_tokens = list(range(num_reqs + 1)) return LogprobsLists( logprob_token_ids=np.array(log_token_ids.tolist()), @@ -1264,6 +1307,7 @@ def _sample_from_logits( else: step_rng = self.rng_params_for_sampling + processed_bonus_logits = None if spec_decode_metadata is None: logits = logits.astype(jnp.float32) with self.maybe_forbid_compile: @@ -1281,7 +1325,7 @@ def _sample_from_logits( rejection_rng = step_rng bonus_logits = self._select_from_array_fn( logits, spec_decode_metadata.bonus_logits_indices) - bonus_token_ids, _ = sample( + bonus_token_ids, processed_bonus_logits = sample( bonus_rng, self.mesh, bonus_logits, @@ -1307,11 +1351,28 @@ def _sample_from_logits( if full_logits is not None: full_logits = full_logits.astype(jnp.float32) with self.maybe_forbid_compile: - if tpu_sampling_metadata.logprobs: - logits = processed_logits if self.model_config.logprobs_mode == "processed_logprobs" else logits + if spec_decode_metadata is not None: + with jax.set_mesh(self.mesh): + if (self.model_config.logprobs_mode + == "processed_logprobs" + and tpu_sampling_metadata.do_sampling): + extended_logits = process_and_extend_logits( + self.mesh, target_logits, + processed_bonus_logits, spec_decode_metadata, + tpu_sampling_metadata) + else: + extended_logits = extend_logits_simple( + target_logits, bonus_logits, self.mesh) + + logprobs_logits = extended_logits + else: + logprobs_logits = (processed_logits + if self.model_config.logprobs_mode + == "processed_logprobs" else logits) logprobs = compute_and_gather_logprobs( - logits, next_tokens, self.model_config.max_logprobs) + logprobs_logits, next_tokens, + self.model_config.max_logprobs) logprobs = _jax_logprobs_copy_to_host_async(logprobs) else: logprobs = None @@ -1480,7 +1541,11 @@ def _sample_from_logits( if logprobs is not None: # Use materialize to ensure logprobs are ready on host when we return async results logprobs_lists = _jax_logprobs_materialize( - logprobs, logits_indices_selector) + logprobs, + logits_indices_selector, + spec_decode_metadata=spec_decode_metadata, + runner=self, + num_reqs=num_reqs) else: logprobs_lists = None @@ -2176,6 +2241,7 @@ def build_attn(block_tables: jax.Array | None) -> AttentionMetadata: all_token_indices_to_substitute) token_in_tpu_pre_next_tokens_indices = np.array( all_pre_next_tokens_indices) + input_ids = self._apply_async_token_substitution( input_ids, next_tokens, token_in_tpu_cur_input_indices, token_in_tpu_pre_next_tokens_indices) diff --git a/tpu_inference/runner/utils.py b/tpu_inference/runner/utils.py index 063d632825..3295da35d6 100644 --- a/tpu_inference/runner/utils.py +++ b/tpu_inference/runner/utils.py @@ -281,7 +281,6 @@ def wrapper(*args, **kwargs): info_after = original_cached_func.cache_info() misses_after = info_after.misses - # Check if a cache miss occurred if misses_after > misses_before: raise RuntimeError(self.message) diff --git a/tpu_inference/spec_decode/jax/utils.py b/tpu_inference/spec_decode/jax/utils.py index ce1bf341da..1fe3c501ee 100644 --- a/tpu_inference/spec_decode/jax/utils.py +++ b/tpu_inference/spec_decode/jax/utils.py @@ -14,9 +14,13 @@ import jax import jax.numpy as jnp +import numpy as np from jax.sharding import PartitionSpec from tpu_inference.layers.common.sharding import ShardingAxisName +from tpu_inference.layers.jax.sample.sampling import _apply_sampling_transforms +from tpu_inference.layers.jax.sample.sampling_metadata import \ + TPUSupportedSamplingMetadata from tpu_inference.runner.utils import SpecDecodeMetadata PLACEHOLDER_TOKEN_ID = -1 @@ -97,3 +101,224 @@ def concat_last_sampled_tokens_and_draft_tokens(last_sampled_tokens, draft_tokens): return jnp.concat([last_sampled_tokens[:, None], draft_tokens], axis=1).reshape(-1) + + +def filter_speculative_logprobs( + log_token_ids: np.ndarray, + logprobs_arr: np.ndarray, + selected_token_ranks: np.ndarray, + spec_decode_metadata: SpecDecodeMetadata, + vocab_size: int, + dp_size: int, + num_reqs: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, list[int]]: + """Filters and reorganizes logprobs for speculative decoding. + + This function extracts logprobs for accepted draft tokens and the bonus token + for each request, and returns them in a flattened format along with cumulative + counts. + + The input arrays (log_token_ids, logprobs_arr, selected_token_ranks) have a flat + layout structured per DP rank. For each rank, the layout is: + [draft_tokens (padded_tokens_length), bonus_tokens (padded_num_seqs_per_rank)] + """ + INVALID_TOKEN_ID = -1 + + # final_logits_indices contains indices for all draft tokens across all ranks. + # We split it by dp_size to get the draft block length per rank. + padded_tokens_length = spec_decode_metadata.final_logits_indices.shape[0] + assert padded_tokens_length % dp_size == 0 + padded_tokens_length = padded_tokens_length // dp_size + + # Total output size per rank (draft block + bonus block) + per_rank_output_size = log_token_ids.shape[0] // dp_size + + # Lists to collect filtered outputs per request (indexed by original request ID) + req_logprob_token_ids = [[] for _ in range(num_reqs)] + req_logprobs = [[] for _ in range(num_reqs)] + req_sampled_token_ranks = [[] for _ in range(num_reqs)] + + # Process outputs rank by rank + for rank in range(dp_size): + rank_offset = rank * per_rank_output_size + + # Slice the draft (main) tokens portion for this rank + main_token_ids = log_token_ids[rank_offset:rank_offset + + padded_tokens_length] + main_logprobs = logprobs_arr[rank_offset:rank_offset + + padded_tokens_length] + main_ranks = selected_token_ranks[rank_offset:rank_offset + + padded_tokens_length] + + # Slice the bonus tokens portion for this rank (comes after draft tokens) + bonus_token_ids = log_token_ids[rank_offset + + padded_tokens_length:rank_offset + + per_rank_output_size] + bonus_logprobs = logprobs_arr[rank_offset + + padded_tokens_length:rank_offset + + per_rank_output_size] + bonus_ranks = selected_token_ranks[rank_offset + + padded_tokens_length:rank_offset + + per_rank_output_size] + + # Number of sequences (requests) allocated for this rank + padded_num_seqs_per_rank = bonus_token_ids.shape[0] + # Actual draft lengths for requests in this rank + cur_rank_num_draft_tokens = spec_decode_metadata.draft_lengths_cpu[ + rank * padded_num_seqs_per_rank:(rank + 1) * + padded_num_seqs_per_rank] + + start_idx = 0 + # Map from rank-local sequence index to global request index + req_indices = spec_decode_metadata.req_indices_dp[rank] + for i, req_idx in enumerate(req_indices): + # Skip padding requests + if req_idx >= num_reqs: + continue + + seq_length = int(cur_rank_num_draft_tokens[i]) + end_idx = start_idx + seq_length + + # Get draft logprobs for this specific sequence + seq_main_token_ids = main_token_ids[start_idx:end_idx] + seq_main_logprobs = main_logprobs[start_idx:end_idx] + seq_main_ranks = main_ranks[start_idx:end_idx] + + # Filter out invalid/padding tokens. + # During speculative execution, some draft positions might be filled with + # placeholders if they are not used or if they are rejected. + valid_mask = (seq_main_token_ids[:, 0] != INVALID_TOKEN_ID) & ( + seq_main_token_ids[:, 0] < vocab_size) + + valid_token_ids = seq_main_token_ids[valid_mask] + valid_logprobs = seq_main_logprobs[valid_mask] + valid_ranks = seq_main_ranks[valid_mask] + + # Check if the bonus token is valid (not a placeholder) + bonus_token = bonus_token_ids[i, 0] + if bonus_token != INVALID_TOKEN_ID and bonus_token < vocab_size: + # Append bonus token logprobs to the valid draft logprobs + valid_token_ids = np.concatenate( + [valid_token_ids, [bonus_token_ids[i]]], axis=0) + valid_logprobs = np.concatenate( + [valid_logprobs, [bonus_logprobs[i]]], axis=0) + valid_ranks = np.concatenate([valid_ranks, [bonus_ranks[i]]], + axis=0) + + # Store the filtered logprobs for this request + req_logprob_token_ids[req_idx] = valid_token_ids + req_logprobs[req_idx] = valid_logprobs + req_sampled_token_ranks[req_idx] = valid_ranks + start_idx = end_idx + + # Flatten the collected lists back to 2D/1D arrays to match the expected output format. + # We also compute cumulative sum of generated tokens (cu_num_generated_tokens). + flat_token_ids = [] + flat_logprobs = [] + flat_ranks = [] + + cu_num_generated_tokens = [0] + current_cu = 0 + + for r in range(num_reqs): + num_gen = len(req_logprob_token_ids[r]) + current_cu += num_gen + cu_num_generated_tokens.append(current_cu) + + if num_gen > 0: + flat_token_ids.append(req_logprob_token_ids[r]) + flat_logprobs.append(req_logprobs[r]) + flat_ranks.append(req_sampled_token_ranks[r]) + + if flat_token_ids: + filtered_token_ids = np.concatenate(flat_token_ids, axis=0) + filtered_logprobs = np.concatenate(flat_logprobs, axis=0) + filtered_ranks = np.concatenate(flat_ranks, axis=0) + else: + # Return empty arrays with correct shape/dtype if no tokens were generated + filtered_token_ids = np.empty((0, log_token_ids.shape[1]), + dtype=log_token_ids.dtype) + filtered_logprobs = np.empty((0, logprobs_arr.shape[1]), + dtype=logprobs_arr.dtype) + filtered_ranks = np.empty((0, ), dtype=selected_token_ranks.dtype) + + return (filtered_token_ids, filtered_logprobs, filtered_ranks, + cu_num_generated_tokens) + + +@jax.jit(static_argnames=["mesh"]) +def extend_logits_simple( + target_logits: jax.Array, + bonus_logits: jax.Array, + mesh: jax.sharding.Mesh, +) -> jax.Array: + """Concatenates target and bonus logits along the first axis.""" + + def concat_fn(x, y): + return jnp.concatenate([x, y], axis=0) + + logits_spec = PartitionSpec(ShardingAxisName.ATTN_DATA, None) + return jax.shard_map( + concat_fn, + mesh=mesh, + in_specs=(logits_spec, logits_spec), + out_specs=logits_spec, + )(target_logits.astype(jnp.float32), bonus_logits.astype(jnp.float32)) + + +@jax.jit(static_argnames=["mesh"]) +def process_and_extend_logits( + mesh: jax.sharding.Mesh, + target_logits: jax.Array, + processed_bonus_logits: jax.Array, + spec_decode_metadata: SpecDecodeMetadata, + tpu_sampling_metadata: TPUSupportedSamplingMetadata, +) -> jax.Array: + """Processes target logits and concatenates them with processed bonus logits.""" + # target_logits: [L, vocab] + # processed_bonus_logits: [R, vocab] + # draft_lengths: [B] + target_logits = target_logits.astype(jnp.float32) + + def local_fn(local_target, local_bonus, local_draft_lengths, local_temp, + local_top_k, local_top_p): + segment_ids = jnp.repeat( + jnp.arange(local_draft_lengths.shape[0]), + local_draft_lengths, + total_repeat_length=local_target.shape[0], + ) + temp = local_temp[segment_ids] + top_k = local_top_k[segment_ids] + top_p = local_top_p[segment_ids] + + local_meta = TPUSupportedSamplingMetadata( + temperature=temp, + top_k=top_k, + top_p=top_p, + do_sampling=tpu_sampling_metadata.do_sampling, + logprobs=tpu_sampling_metadata.logprobs, + ) + processed_target = _apply_sampling_transforms(local_target, local_meta) + return jnp.concatenate([processed_target, local_bonus], axis=0) + + return jax.shard_map( + local_fn, + mesh=mesh, + in_specs=( + PartitionSpec(ShardingAxisName.ATTN_DATA), # target_logits + PartitionSpec( + ShardingAxisName.ATTN_DATA), # processed_bonus_logits + PartitionSpec(ShardingAxisName.ATTN_DATA), # draft_lengths + PartitionSpec(ShardingAxisName.ATTN_DATA), # temperature + PartitionSpec(ShardingAxisName.ATTN_DATA), # top_k + PartitionSpec(ShardingAxisName.ATTN_DATA), # top_p + ), + out_specs=PartitionSpec(ShardingAxisName.ATTN_DATA), + )( + target_logits, + processed_bonus_logits, + spec_decode_metadata.draft_lengths, + tpu_sampling_metadata.temperature, + tpu_sampling_metadata.top_k, + tpu_sampling_metadata.top_p, + ) From 6a8f0a4f055459345bb47f2cea531617b0623a5f Mon Sep 17 00:00:00 2001 From: muskansh-google Date: Fri, 5 Jun 2026 10:44:15 +0530 Subject: [PATCH 10/37] Add qwen3-vl run in example script (#2763) --- examples/multi_modal_inference.py | 86 +++++++++++++++++++++++-------- 1 file changed, 65 insertions(+), 21 deletions(-) diff --git a/examples/multi_modal_inference.py b/examples/multi_modal_inference.py index 7dc96d01bc..1ce2327528 100644 --- a/examples/multi_modal_inference.py +++ b/examples/multi_modal_inference.py @@ -12,6 +12,12 @@ --model Qwen/Qwen2.5-VL-3B-Instruct \ --tensor-parallel-size 1 \ --num-prompts 1 + +Example command to test multiple images +python examples/multi_modal_inference.py \ + --model Qwen/Qwen3-VL-8B-Instruct \ + --test-multi-image \ + --max-model-len 8192 """ from contextlib import contextmanager @@ -30,15 +36,14 @@ class ModelRequestData(NamedTuple): stop_token_ids: Optional[list[int]] = None -# Currently Qwen2.5-VL is the only supported multi-modal -# Qwen2.5-VL -def run_qwen2_5_vl(questions: list[str], modality: str, - args) -> ModelRequestData: +# Currently Qwen2.5-VL and Qwen3-VL are supported +def run_qwen_vl(questions: list[str], modality: str, args) -> ModelRequestData: engine_args = EngineArgs( model=args.model, max_model_len=args.max_model_len, tensor_parallel_size=args.tensor_parallel_size, gpu_memory_utilization=args.gpu_memory_utilization, + enable_chunked_prefill=False, max_num_seqs=5, mm_processor_kwargs={ "size": { @@ -47,7 +52,7 @@ def run_qwen2_5_vl(questions: list[str], modality: str, }, "fps": 1, }, - limit_mm_per_prompt={modality: 1}, + limit_mm_per_prompt={modality: 2 if args.test_multi_image else 1}, ) if modality == "image": @@ -55,12 +60,14 @@ def run_qwen2_5_vl(questions: list[str], modality: str, elif modality == "video": placeholder = "<|video_pad|>" - prompts = [ - ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" - f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>" - f"{question}<|im_end|>\n" - "<|im_start|>assistant\n") for question in questions - ] + placeholder_full = f"<|vision_start|>{placeholder}<|vision_end|>" + if args.test_multi_image: + placeholder_full += f"<|vision_start|>{placeholder}<|vision_end|>" + + prompts = [("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" + f"<|im_start|>user\n{placeholder_full}" + f"{question}<|im_end|>\n" + "<|im_start|>assistant\n") for question in questions] return ModelRequestData( engine_args=engine_args, @@ -68,9 +75,13 @@ def run_qwen2_5_vl(questions: list[str], modality: str, ) -model_example_map = { - "qwen2_5_vl": run_qwen2_5_vl, -} +def get_model_runner(model_name: str): + """Returns the appropriate run function for the given model.""" + if "qwen" in model_name.lower(): + return run_qwen_vl + + raise ValueError(f"Unsupported model: {model_name}. " + "Please add support for this model in get_model_runner.") def get_multi_modal_input(args): @@ -96,6 +107,30 @@ def get_multi_modal_input(args): "questions": img_questions, } + msg = f"Modality {args.modality} is not supported." + raise ValueError(msg) + + +def get_multi_modal_input_multi(args): + """ + Returns multiple images for testing compare. + """ + if args.modality == "image": + image1 = convert_image_mode( + ImageAsset("cherry_blossom").pil_image, "RGB") + image2 = convert_image_mode(ImageAsset("stop_sign").pil_image, "RGB") + images = [image1, image2] + img_questions = [ + "What are shown in these two images? Compare them.", + "Describe the content of both images and how they differ.", + "What's in the first image vs the second image?", + ] + + return { + "data": images, + "questions": img_questions, + } + # NOTE: not used for now, saved for future reference # if args.modality == "video": # # Input video and question @@ -152,7 +187,7 @@ def parse_args(): parser.add_argument( "--gpu-memory-utilization", type=float, - default=0.5, + default=0.85, help="GPU memory utilization", ) @@ -190,10 +225,11 @@ def parse_args(): default=None, help="Set the seed when initializing `vllm.LLM`.", ) + parser.add_argument( - "--disable-mm-preprocessor-cache", + "--test-multi-image", action="store_true", - help="If True, disables caching of multi-modal preprocessor/mapper.", + help="If set, run the multiple images test (Option B).", ) parser.add_argument( @@ -215,12 +251,16 @@ def parse_args(): def main(args): modality = args.modality - mm_input = get_multi_modal_input(args) + if args.test_multi_image: + mm_input = get_multi_modal_input_multi(args) + else: + mm_input = get_multi_modal_input(args) data = mm_input["data"] questions = mm_input["questions"] - # NOTE: Currently, only Qwen2.5-VL is supported. If later we want to support a model with new chat template, we may need to change this - req_data = model_example_map["qwen2_5_vl"](questions, modality, args) + # NOTE: Currently, only Qwen2.5-VL and Qwen3-VL is supported. If later we want to support a model with new chat template, we may need to change this + model_key = args.model + req_data = get_model_runner(model_key)(questions, modality, args) # Disable other modalities to save memory # Initial all modalities to be 0s and add the specifc modality limit later accordingly @@ -235,6 +275,10 @@ def main(args): if engine_args.get("compilation_config") is None: engine_args["compilation_config"] = {} engine_args["compilation_config"]["cudagraph_capture_sizes"] = [] + # Fix for pydantic validation error in recent vLLM versions + engine_args["compilation_config"]["pass_config"] = { + "fuse_minimax_qk_norm": False + } llm = LLM(**engine_args) @@ -245,7 +289,7 @@ def main(args): # We set temperature to 0.2 so that outputs can be different # even when all prompts are identical when running batch inference. sampling_params = SamplingParams(temperature=0.2, - max_tokens=64, + max_tokens=1024, stop_token_ids=req_data.stop_token_ids) assert args.num_prompts > 0 From 7612c22cd84849140580188434c1f38d8fc24f5b Mon Sep 17 00:00:00 2001 From: Lumosis <30372757+Lumosis@users.noreply.github.com> Date: Thu, 4 Jun 2026 22:33:25 -0700 Subject: [PATCH 11/37] [MTP] Add MTP to support matrics (#2826) --- support_matrices/nightly/v7x/vllm/feature_support_matrix.csv | 1 + support_matrices/release/v7x/vllm/feature_support_matrix.csv | 1 + 2 files changed, 2 insertions(+) diff --git a/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv b/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv index a9124ad39b..c79974d046 100644 --- a/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv +++ b/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv @@ -9,6 +9,7 @@ Feature,CorrectnessTest,PerformanceTest "Single Program Multi Data",✅ Passing,✅ Passing "Speculative Decoding: Eagle3",✅ Passing,✅ Passing "Speculative Decoding: Ngram",✅ Passing,✅ Passing +"Speculative Decoding: MTP",✅ Passing,✅ Passing "Step Pooling (Embedding)",✅ Passing,❓ Untested "async scheduler",✅ Passing,✅ Passing "hybrid kv cache",✅ Passing,❓ Untested diff --git a/support_matrices/release/v7x/vllm/feature_support_matrix.csv b/support_matrices/release/v7x/vllm/feature_support_matrix.csv index 97b24d8eb8..4ab7cfde4f 100644 --- a/support_matrices/release/v7x/vllm/feature_support_matrix.csv +++ b/support_matrices/release/v7x/vllm/feature_support_matrix.csv @@ -10,6 +10,7 @@ Feature,CorrectnessTest,PerformanceTest "Single-Host-P-D-disaggregation",✅ Passing,❓ Untested "Speculative Decoding: Eagle3",❌ Failing,❓ Untested "Speculative Decoding: Ngram",✅ Passing,✅ Passing +"Speculative Decoding: MTP",✅ Passing,✅ Passing "async scheduler",✅ Passing,✅ Passing "hybrid kv cache",✅ Passing,❓ Untested "multi-host",✅ Passing,❓ Untested From 25e5d14c1a5beba3681dc52b3941bbf7699b6ce9 Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Fri, 5 Jun 2026 14:17:45 +0000 Subject: [PATCH 12/37] [skip ci] Update nightly support matrices (v6e/v7x) Signed-off-by: Buildkite Bot --- .../nightly/v6e/default/feature_support_matrix.csv | 6 +++--- .../nightly/v6e/default/parallelism_support_matrix.csv | 2 +- support_matrices/nightly/v6e/default/rl_support_matrix.csv | 2 +- .../nightly/v7x/default/parallelism_support_matrix.csv | 2 +- support_matrices/nightly/v7x/default/rl_support_matrix.csv | 2 +- 5 files changed, 7 insertions(+), 7 deletions(-) diff --git a/support_matrices/nightly/v6e/default/feature_support_matrix.csv b/support_matrices/nightly/v6e/default/feature_support_matrix.csv index 57a12bf226..d085f597b5 100644 --- a/support_matrices/nightly/v6e/default/feature_support_matrix.csv +++ b/support_matrices/nightly/v6e/default/feature_support_matrix.csv @@ -1,7 +1,7 @@ Feature,CorrectnessTest,PerformanceTest "Chunked Prefill",✅ Passing,✅ Passing -"DCN-based P/D disaggregation",❌ Failing,✅ Passing -"KV Cache Offload",✅ Passing,✅ Passing +"DCN-based P/D disaggregation",✅ Passing,✅ Passing +"KV Cache Offload",❌ Failing,❓ Untested "LoRA_Torch",✅ Passing,✅ Passing "Multimodal Inputs",✅ Passing,✅ Passing "Out-of-tree model support",✅ Passing,✅ Passing @@ -10,7 +10,7 @@ Feature,CorrectnessTest,PerformanceTest "Speculative Decoding: Eagle3",✅ Passing,✅ Passing "Speculative Decoding: Ngram",✅ Passing,✅ Passing "Step Pooling (Embedding)",✅ Passing,❓ Untested -"async scheduler",❌ Failing,❓ Untested +"async scheduler",✅ Passing,✅ Passing "hybrid kv cache",✅ Passing,❓ Untested "multi-host",❓ Untested,❓ Untested "runai_model_streamer_loader",✅ Passing,❓ Untested diff --git a/support_matrices/nightly/v6e/default/parallelism_support_matrix.csv b/support_matrices/nightly/v6e/default/parallelism_support_matrix.csv index 5eeba3be1f..a6cc9507ed 100644 --- a/support_matrices/nightly/v6e/default/parallelism_support_matrix.csv +++ b/support_matrices/nightly/v6e/default/parallelism_support_matrix.csv @@ -3,5 +3,5 @@ Feature,Single-Host CorrectnessTest,Single-Host PerformanceTest,Multi-Host Corre "DP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "EP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "PP",✅ Passing,✅ Passing,✅ Passing,✅ Passing -"SP",❌ Failing,❓ Untested,❓ Untested,❓ Untested +"SP",✅ Passing,❓ Untested,❓ Untested,❓ Untested "TP",✅ Passing,✅ Passing,❓ Untested,❓ Untested diff --git a/support_matrices/nightly/v6e/default/rl_support_matrix.csv b/support_matrices/nightly/v6e/default/rl_support_matrix.csv index 22a22d20a6..9174f88238 100644 --- a/support_matrices/nightly/v6e/default/rl_support_matrix.csv +++ b/support_matrices/nightly/v6e/default/rl_support_matrix.csv @@ -1,6 +1,6 @@ Feature,CorrectnessTest,PerformanceTest "Group sampling prefix cache",✅ Passing,✅ Passing -"MoE Expert IDs",❌ Failing,❓ Untested +"MoE Expert IDs",✅ Passing,❓ Untested "async_logprobs",✅ Passing,❓ Untested "processed_logprobs",❌ Failing,❓ Untested "rl_integration",✅ Passing,✅ Passing diff --git a/support_matrices/nightly/v7x/default/parallelism_support_matrix.csv b/support_matrices/nightly/v7x/default/parallelism_support_matrix.csv index 5eeba3be1f..a6cc9507ed 100644 --- a/support_matrices/nightly/v7x/default/parallelism_support_matrix.csv +++ b/support_matrices/nightly/v7x/default/parallelism_support_matrix.csv @@ -3,5 +3,5 @@ Feature,Single-Host CorrectnessTest,Single-Host PerformanceTest,Multi-Host Corre "DP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "EP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "PP",✅ Passing,✅ Passing,✅ Passing,✅ Passing -"SP",❌ Failing,❓ Untested,❓ Untested,❓ Untested +"SP",✅ Passing,❓ Untested,❓ Untested,❓ Untested "TP",✅ Passing,✅ Passing,❓ Untested,❓ Untested diff --git a/support_matrices/nightly/v7x/default/rl_support_matrix.csv b/support_matrices/nightly/v7x/default/rl_support_matrix.csv index 22a22d20a6..9174f88238 100644 --- a/support_matrices/nightly/v7x/default/rl_support_matrix.csv +++ b/support_matrices/nightly/v7x/default/rl_support_matrix.csv @@ -1,6 +1,6 @@ Feature,CorrectnessTest,PerformanceTest "Group sampling prefix cache",✅ Passing,✅ Passing -"MoE Expert IDs",❌ Failing,❓ Untested +"MoE Expert IDs",✅ Passing,❓ Untested "async_logprobs",✅ Passing,❓ Untested "processed_logprobs",❌ Failing,❓ Untested "rl_integration",✅ Passing,✅ Passing From 08c01e013c87edbd0ddd9f0fa1d9348b2d3b8b3e Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Fri, 5 Jun 2026 14:59:33 +0000 Subject: [PATCH 13/37] [skip ci] Update nightly support matrices for flax_nnx (v6e/v7x) Signed-off-by: Buildkite Bot --- .../nightly/v6e/flax_nnx/parallelism_support_matrix.csv | 2 +- .../nightly/v7x/flax_nnx/parallelism_support_matrix.csv | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/support_matrices/nightly/v6e/flax_nnx/parallelism_support_matrix.csv b/support_matrices/nightly/v6e/flax_nnx/parallelism_support_matrix.csv index 5eeba3be1f..a6cc9507ed 100644 --- a/support_matrices/nightly/v6e/flax_nnx/parallelism_support_matrix.csv +++ b/support_matrices/nightly/v6e/flax_nnx/parallelism_support_matrix.csv @@ -3,5 +3,5 @@ Feature,Single-Host CorrectnessTest,Single-Host PerformanceTest,Multi-Host Corre "DP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "EP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "PP",✅ Passing,✅ Passing,✅ Passing,✅ Passing -"SP",❌ Failing,❓ Untested,❓ Untested,❓ Untested +"SP",✅ Passing,❓ Untested,❓ Untested,❓ Untested "TP",✅ Passing,✅ Passing,❓ Untested,❓ Untested diff --git a/support_matrices/nightly/v7x/flax_nnx/parallelism_support_matrix.csv b/support_matrices/nightly/v7x/flax_nnx/parallelism_support_matrix.csv index 5eeba3be1f..a6cc9507ed 100644 --- a/support_matrices/nightly/v7x/flax_nnx/parallelism_support_matrix.csv +++ b/support_matrices/nightly/v7x/flax_nnx/parallelism_support_matrix.csv @@ -3,5 +3,5 @@ Feature,Single-Host CorrectnessTest,Single-Host PerformanceTest,Multi-Host Corre "DP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "EP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "PP",✅ Passing,✅ Passing,✅ Passing,✅ Passing -"SP",❌ Failing,❓ Untested,❓ Untested,❓ Untested +"SP",✅ Passing,❓ Untested,❓ Untested,❓ Untested "TP",✅ Passing,✅ Passing,❓ Untested,❓ Untested From 899ebc6878dd9c9398327f130ce3da25f591ff63 Mon Sep 17 00:00:00 2001 From: Amanda Liang Date: Fri, 5 Jun 2026 09:15:13 -0700 Subject: [PATCH 14/37] Customize MLA kernel tuner keys (#2824) Signed-off-by: Amanda Liang --- tools/kernel/tuner/v1/mla_kernel_tuner.py | 41 +++++++++++++++++------ 1 file changed, 30 insertions(+), 11 deletions(-) diff --git a/tools/kernel/tuner/v1/mla_kernel_tuner.py b/tools/kernel/tuner/v1/mla_kernel_tuner.py index 4faf76fd16..3ad4fee31c 100644 --- a/tools/kernel/tuner/v1/mla_kernel_tuner.py +++ b/tools/kernel/tuner/v1/mla_kernel_tuner.py @@ -18,6 +18,7 @@ import jax import jax.numpy as jnp import numpy as np +from absl import flags from vllm.utils.math_utils import cdiv from tools.kernel.tuner.v1.common.kernel_tuner_base import (KernelTunerBase, @@ -32,6 +33,21 @@ logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) +flags.DEFINE_integer("mla_total_num_pages", 1506, + "Total number of pages in the cache.") +flags.DEFINE_integer("mla_page_size_per_kv_packing", 256, + "Page size per KV packing.") +flags.DEFINE_integer("mla_kv_packing", 4, "Packing factor for KV.") +flags.DEFINE_integer("mla_max_num_seqs", 160, + "Maximum number of sequences in the batch.") +flags.DEFINE_integer("mla_pages_per_seq", 9, "Number of pages per sequence.") +flags.DEFINE_integer("mla_actual_num_q_heads", 128, + "Actual number of Q heads.") +flags.DEFINE_integer("mla_actual_lkv_dim", 512, "Actual NOPE head dimension.") +flags.DEFINE_integer("mla_actual_r_dim", 64, "Actual ROPE head dimension.") +flags.DEFINE_string("mla_kv_dtype", "float8_e4m3fn", "KV cache data type.") +flags.DEFINE_string("mla_q_dtype", "float8_e4m3fn", "Q activation dtype.") + def _generate_mla_inputs( seq_lens, # List[(q_len, kv_len)] @@ -153,20 +169,23 @@ def generate_cases(self) -> list[TuningCase]: # The tuningKey and tunableParams in this list is constructed based on the real kernel execution logs during benchmarking. # The tunableParams currently represent the baseline configuration from the attention_interface.py. # From the log, only the max_num_tokens is different across different cases - for max_num_tokens in [4, 8, 16, 32, 64, 128, 160, 256, 512]: + for max_num_tokens in [ + 4, 8, 16, 32, 64, 128, 160, 256, 512, 1024, 2048 + ]: tuning_set_from_log.append([ TuningKey( max_num_tokens=max_num_tokens, - actual_num_q_heads=128, - actual_lkv_dim=512, - actual_r_dim=64, - kv_dtype="float8_e4m3fn", - q_dtype="float8_e4m3fn", - total_num_pages=1506, - page_size_per_kv_packing=256, - kv_packing=4, - max_num_seqs=160, - pages_per_seq=9, + actual_num_q_heads=flags.FLAGS.mla_actual_num_q_heads, + actual_lkv_dim=flags.FLAGS.mla_actual_lkv_dim, + actual_r_dim=flags.FLAGS.mla_actual_r_dim, + kv_dtype=flags.FLAGS.mla_kv_dtype, + q_dtype=flags.FLAGS.mla_q_dtype, + total_num_pages=flags.FLAGS.mla_total_num_pages, + page_size_per_kv_packing=flags.FLAGS. + mla_page_size_per_kv_packing, + kv_packing=flags.FLAGS.mla_kv_packing, + max_num_seqs=flags.FLAGS.mla_max_num_seqs, + pages_per_seq=flags.FLAGS.mla_pages_per_seq, s_dtype="bfloat16", case="batched_decode", soft_cap=None, From 9f5ae248aea638c31365311a44ae5bc7dd18aafd Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Fri, 5 Jun 2026 16:19:02 +0000 Subject: [PATCH 15/37] [skip ci] Update nightly support matrices for vllm (v6e/v7x) Signed-off-by: Buildkite Bot --- .../nightly/v6e/vllm/parallelism_support_matrix.csv | 2 +- support_matrices/nightly/v7x/vllm/feature_support_matrix.csv | 5 ++--- .../nightly/v7x/vllm/parallelism_support_matrix.csv | 2 +- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv b/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv index 5eeba3be1f..a6cc9507ed 100644 --- a/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv +++ b/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv @@ -3,5 +3,5 @@ Feature,Single-Host CorrectnessTest,Single-Host PerformanceTest,Multi-Host Corre "DP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "EP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "PP",✅ Passing,✅ Passing,✅ Passing,✅ Passing -"SP",❌ Failing,❓ Untested,❓ Untested,❓ Untested +"SP",✅ Passing,❓ Untested,❓ Untested,❓ Untested "TP",✅ Passing,✅ Passing,❓ Untested,❓ Untested diff --git a/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv b/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv index c79974d046..21b2ea4ef7 100644 --- a/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv +++ b/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv @@ -3,17 +3,16 @@ Feature,CorrectnessTest,PerformanceTest "DCN-based P/D disaggregation",✅ Passing,✅ Passing "KV Cache Offload",✅ Passing,✅ Passing "LoRA_Torch",✅ Passing,✅ Passing -"Multimodal Inputs",✅ Passing,✅ Passing +"Multimodal Inputs",❌ Failing,❓ Untested "Out-of-tree model support",✅ Passing,✅ Passing "Prefix Caching",✅ Passing,✅ Passing "Single Program Multi Data",✅ Passing,✅ Passing "Speculative Decoding: Eagle3",✅ Passing,✅ Passing "Speculative Decoding: Ngram",✅ Passing,✅ Passing -"Speculative Decoding: MTP",✅ Passing,✅ Passing "Step Pooling (Embedding)",✅ Passing,❓ Untested "async scheduler",✅ Passing,✅ Passing "hybrid kv cache",✅ Passing,❓ Untested -"multi-host",❌ Failing,❓ Untested +"multi-host",✅ Passing,❓ Untested "runai_model_streamer_loader",✅ Passing,❓ Untested "sampling_params",✅ Passing,❓ Untested "structured_decoding",✅ Passing,❓ Untested diff --git a/support_matrices/nightly/v7x/vllm/parallelism_support_matrix.csv b/support_matrices/nightly/v7x/vllm/parallelism_support_matrix.csv index 5eeba3be1f..a6cc9507ed 100644 --- a/support_matrices/nightly/v7x/vllm/parallelism_support_matrix.csv +++ b/support_matrices/nightly/v7x/vllm/parallelism_support_matrix.csv @@ -3,5 +3,5 @@ Feature,Single-Host CorrectnessTest,Single-Host PerformanceTest,Multi-Host Corre "DP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "EP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "PP",✅ Passing,✅ Passing,✅ Passing,✅ Passing -"SP",❌ Failing,❓ Untested,❓ Untested,❓ Untested +"SP",✅ Passing,❓ Untested,❓ Untested,❓ Untested "TP",✅ Passing,✅ Passing,❓ Untested,❓ Untested From 84c992a1b207c64f0c5e87b97a9c8a44e82b93ac Mon Sep 17 00:00:00 2001 From: Yiwei Wang Date: Fri, 5 Jun 2026 11:30:06 -0700 Subject: [PATCH 16/37] Remove the logic to force change the compilation_config (#2821) Signed-off-by: Yiwei Wang --- tests/platforms/test_tpu_platform.py | 7 ++----- tpu_inference/platforms/tpu_platform.py | 15 ++------------- 2 files changed, 4 insertions(+), 18 deletions(-) diff --git a/tests/platforms/test_tpu_platform.py b/tests/platforms/test_tpu_platform.py index e25241a1d3..df2f658544 100644 --- a/tests/platforms/test_tpu_platform.py +++ b/tests/platforms/test_tpu_platform.py @@ -18,7 +18,7 @@ import jax.numpy as jnp import pytest import torch -from vllm.config import CacheConfig, CompilationMode, ModelConfig, VllmConfig +from vllm.config import CacheConfig, ModelConfig, VllmConfig from tpu_inference.platforms.tpu_platform import TpuPlatform @@ -48,8 +48,7 @@ def vllm_config(self): vllm_config.parallel_config = MagicMock() vllm_config.parallel_config.data_parallel_size = 1 vllm_config.sharding_config = MagicMock() - vllm_config.compilation_config = MagicMock(mode="dynamo_trace_once", - backend="openxla") + vllm_config.compilation_config = MagicMock(backend="eager") vllm_config.kv_transfer_config = None return vllm_config @@ -188,8 +187,6 @@ def test_check_and_update_config_single_host_uni(self, mock_update, TpuPlatform.check_and_update_config(vllm_config) - assert vllm_config.compilation_config.mode == CompilationMode.DYNAMO_TRACE_ONCE - assert vllm_config.compilation_config.backend == "openxla" assert vllm_config.parallel_config.distributed_executor_backend == "uni" assert vllm_config.scheduler_config.disable_chunked_mm_input is False diff --git a/tpu_inference/platforms/tpu_platform.py b/tpu_inference/platforms/tpu_platform.py index 22f0fe8523..b5576e8ab2 100644 --- a/tpu_inference/platforms/tpu_platform.py +++ b/tpu_inference/platforms/tpu_platform.py @@ -105,7 +105,8 @@ class TpuPlatform(Platform): dispatch_key: str = "XLA" ray_device_key: str = "TPU" device_control_env_var: str = "TPU_VISIBLE_CHIPS" - simple_compile_backend: str = "openxla" + # Bypass torch.compile; torchax defers all compilation to JAX + simple_compile_backend: str = "eager" supported_quantization: list[str] = [ "compressed-tensors", "awq", "fp8", "gpt_oss_mxfp4", "modelopt_fp4", @@ -259,18 +260,6 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: ) cls._initialize_sharding_config(vllm_config) - from vllm.config import CompilationMode - - compilation_config = vllm_config.compilation_config - - # TPU only supports DYNAMO_TRACE_ONCE compilation level - # NOTE(xiang): the compilation_config is not used by jax. - if compilation_config.mode != CompilationMode.DYNAMO_TRACE_ONCE: - compilation_config.mode = CompilationMode.DYNAMO_TRACE_ONCE - - if compilation_config.backend == "": - compilation_config.backend = "openxla" - cache_config = vllm_config.cache_config # For v0, the default block size is 16. if cache_config and not cache_config.user_specified_block_size: From 57c68ca1ca64570a18bc150dd4bd9c7c0c45a7ec Mon Sep 17 00:00:00 2001 From: Sierra Qian <133469784+sierraisland@users.noreply.github.com> Date: Fri, 5 Jun 2026 14:04:10 -0700 Subject: [PATCH 17/37] [Feature] Speed up compilation via AoT + threadpool (#2818) Signed-off-by: wenxindongwork Signed-off-by: Sierra Qian <133469784+sierraisland@users.noreply.github.com> Co-authored-by: wenxindongwork --- .buildkite/scripts/run_in_docker.sh | 1 + tests/runner/test_tpu_runner_dp.py | 6 +- tpu_inference/envs.py | 4 + tpu_inference/runner/compilation_manager.py | 308 ++++++++++++++------ tpu_inference/runner/tpu_runner.py | 35 +-- tpu_inference/worker/tpu_worker.py | 1 + 6 files changed, 245 insertions(+), 110 deletions(-) diff --git a/.buildkite/scripts/run_in_docker.sh b/.buildkite/scripts/run_in_docker.sh index b6c3969bbd..c240d1c1c0 100755 --- a/.buildkite/scripts/run_in_docker.sh +++ b/.buildkite/scripts/run_in_docker.sh @@ -182,6 +182,7 @@ docker run \ ${USE_V7X8_QUEUE:+-e USE_V7X8_QUEUE="$USE_V7X8_QUEUE"} \ ${MOE_REQUANTIZE_BLOCK_SIZE:+-e MOE_REQUANTIZE_BLOCK_SIZE="$MOE_REQUANTIZE_BLOCK_SIZE"} \ ${MOE_REQUANTIZE_WEIGHT_DTYPE:+-e MOE_REQUANTIZE_WEIGHT_DTYPE="$MOE_REQUANTIZE_WEIGHT_DTYPE"} \ + -e NUM_PRECOMPILE_WORKERS="${NUM_PRECOMPILE_WORKERS:-1}" \ "${BENCHMARK_DOCKER_ARGS[@]}" \ "$FULL_IMAGE_TAG" \ "$@" # Pass all script arguments as the command to run in the container diff --git a/tests/runner/test_tpu_runner_dp.py b/tests/runner/test_tpu_runner_dp.py index 14cf2de6df..fb8a3a8b4d 100644 --- a/tests/runner/test_tpu_runner_dp.py +++ b/tests/runner/test_tpu_runner_dp.py @@ -1198,8 +1198,12 @@ def test_get_intermediate_tensor_spec(self): class TestSamplingMetadataPassthrough: - def test_sample_tokens_passes_sampling_metadata_from_state(self): + @patch('jax.set_mesh') + def test_sample_tokens_passes_sampling_metadata_from_state( + self, mock_set_mesh): """sample_tokens() should pass execute_model_state.sampling_metadata to _sample_from_logits.""" + mock_set_mesh.return_value.__enter__ = MagicMock(return_value=None) + mock_set_mesh.return_value.__exit__ = MagicMock(return_value=False) runner = MagicMock() mock_sampling_metadata = MagicMock() # Set the specific field we want to trace; other fields are auto-mocked. diff --git a/tpu_inference/envs.py b/tpu_inference/envs.py index 8be0ac6472..bfeaf73507 100644 --- a/tpu_inference/envs.py +++ b/tpu_inference/envs.py @@ -57,6 +57,7 @@ MOE_APPROX_TOPK_RECALL_TARGET: float | None = None VLLM_TPU_PATCH_MM_EMBEDDINGS: bool = False ENABLE_RS_KERNEL: bool = False + NUM_PRECOMPILE_WORKERS: int = 1 DP_SCHED_BATCH_PREFILL: bool = False DP_SCHED_BATCH_PREFILL_FLUSH_TIMEOUT_MS: int = 10000 ONEHOT_MOE_PERMUTE_THRESHOLD: int = 0 @@ -369,6 +370,9 @@ def _get_int_list_env() -> list[int]: # Enable hierarchical reduce-scatter kernel for MoE "ENABLE_RS_KERNEL": env_bool("ENABLE_RS_KERNEL", default=False), + # Number of worker threads for parallel XLA precompilation. + "NUM_PRECOMPILE_WORKERS": + lambda: int(os.getenv("NUM_PRECOMPILE_WORKERS") or "1"), # DP scheudler: hold and batch incoming requests (prefills) to # cluster and dispatch prefills together. "DP_SCHED_BATCH_PREFILL": diff --git a/tpu_inference/runner/compilation_manager.py b/tpu_inference/runner/compilation_manager.py index 164f4f9954..9df471c751 100644 --- a/tpu_inference/runner/compilation_manager.py +++ b/tpu_inference/runner/compilation_manager.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools +import threading import time +from concurrent.futures import Future, ThreadPoolExecutor from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple import jax @@ -37,7 +40,7 @@ concat_last_sampled_tokens_and_draft_tokens, extend_logits_simple, extract_last_sampled_tokens, process_and_extend_logits) from tpu_inference.utils import (device_array, get_mesh_shape_product, - to_jax_dtype) + time_function, to_jax_dtype) if TYPE_CHECKING: from tpu_inference.runner.tpu_runner import TPUModelRunner @@ -64,6 +67,29 @@ def __init__(self, runner: "TPUModelRunner"): -1) jax.config.update("jax_persistent_cache_min_compile_time_secs", -1) + # Thread pool for parallel XLA compilation. NUM_PRECOMPILE_WORKERS=1 + # disables the pool and runs compilations sequentially in the main + # thread. + num_workers = envs.NUM_PRECOMPILE_WORKERS + self._prev_stack_size: Optional[int] = None + if num_workers == 1: + self._compile_executor = None + else: + # Pool threads default to the system thread stack size (~8MB on + # Linux), much smaller than the main thread. XLA lowering overflows + # that stack on large graphs. Bump the default stack size. + try: + self._prev_stack_size = threading.stack_size() + threading.stack_size(64 * 1024 * 1024) + except (RuntimeError, ValueError): + self._prev_stack_size = None + logger.info( + "Parallel AOT compilation enabled (NUM_PRECOMPILE_WORKERS=%d).", + num_workers) + self._compile_executor = ThreadPoolExecutor( + max_workers=num_workers, thread_name_prefix="aot_compilation") + self._compile_futures: list[Future] = [] + self._warmup_tasks: list = [] def _create_dummy_tensor(self, shape: Tuple[int, ...], @@ -98,56 +124,161 @@ def _run_compilation(self, fn: Callable, *args, call_kwargs=dict(), + warmup_handler: Optional[Callable] = None, + aot: bool = True, **kwargs) -> None: - logger.info(f"Precompile {name} --> {kwargs}") - start = time.perf_counter() - result = fn(*args, **call_kwargs) - jax.tree.map(lambda r: r.block_until_ready(), result) - end = time.perf_counter() - logger.info("Compilation finished in %.2f [secs].", end - start) + log_name = f"{name} --> {kwargs}" + logger.info(f"Precompile {log_name}") + # Unwrap functools.partial so the underlying jit's static_argnums are + # respected. + while isinstance(fn, functools.partial): + args = fn.args + args + call_kwargs = {**fn.keywords, **call_kwargs} + fn = fn.func + self._warmup_tasks.append( + (name, fn, args, call_kwargs, warmup_handler)) + if not aot or not hasattr(fn, 'lower'): + # Skip AOT when the caller opts out, or when fn is unjitted. + # The warmup pass will run fn() and populate the inner-jit caches. + reason = "aot=False" if not aot else "not a jit" + logger.info( + "AOT lower skipped for %s (%s); will compile in warmup.", name, + reason) + return + try: + lowered = fn.lower(*args, **call_kwargs) + except Exception as e: + # AOT lower not supported here (e.g. a jit whose body contains a + # nested jit with compiler_options). Fall back to warmup-only — the + # warmup pass will trigger inline compile. + logger.info( + "AOT lower skipped for %s (%r); will compile in warmup.", name, + e) + return + + # Compilation is thread-safe + def _compile(lowered, name, mesh): + with jax.set_mesh(mesh): + start = time.perf_counter() + compiled = lowered.compile() + elapsed = time.perf_counter() - start + logger.info("Compilation of %s finished in %.2f [secs].", name, + elapsed) + return compiled + + if self._compile_executor is None: + _compile(lowered, log_name, self.runner.mesh) + else: + future = self._compile_executor.submit(_compile, lowered, log_name, + self.runner.mesh) + self._compile_futures.append(future) + + def _flush_compilations(self) -> None: + """Wait for all currently-pending background compilations and run their + warmups. + """ + futures, self._compile_futures = self._compile_futures, [] + tasks, self._warmup_tasks = self._warmup_tasks, [] + + for fut in futures: + try: + fut.result() + except Exception as e: + raise RuntimeError( + f"Compilation failed: {e}\n" + "Hint: if you are seeing memory errors or stack overflows " + "during parallel precompilation, try lowering " + "NUM_PRECOMPILE_WORKERS (e.g. NUM_PRECOMPILE_WORKERS=1 " + "runs all compilations sequentially in the main thread)." + ) from e + + warmup_start = time.perf_counter() + with jax.set_mesh(self.runner.mesh): + for name, fn, args, call_kwargs, warmup_handler in tasks: + if warmup_handler is not None: + out = warmup_handler(fn, args, call_kwargs) + else: + out = fn(*args, **call_kwargs) + jax.tree.map(lambda r: r.block_until_ready(), out) + warmup_elapsed = time.perf_counter() - warmup_start + if tasks: + logger.info( + "Warm-up call pass finished in %.2f [secs] over %d tasks.", + warmup_elapsed, len(tasks)) + @time_function def capture_model(self) -> None: if envs.SKIP_JAX_PRECOMPILE or self.runner.model_config.enforce_eager: return logger.info("Precompile all the subgraphs with possible input shapes.") compilation_start_time = time.perf_counter() - with self.runner.maybe_setup_dummy_loras( - self.runner.lora_config), jax.set_mesh(self.runner.mesh): - self._precompile_backbone_text_only() - if self.runner.is_multimodal_model: - if self.runner.precompile_vision_encoder_fn is not None: - self.runner.precompile_vision_encoder_fn( - self._run_compilation, ) - self._precompile_input_embeddings_merger() - self._precompile_backbone_with_inputs_embeds() - if self.runner.scheduler_config.async_scheduling: - self._precompile_substitute_placeholder_token() + try: + with self.runner.maybe_setup_dummy_loras( + self.runner.lora_config), jax.set_mesh(self.runner.mesh): + self._precompile_backbone_text_only() + self._flush_compilations() + if self.runner.is_multimodal_model: + if self.runner.precompile_vision_encoder_fn is not None: + self.runner.precompile_vision_encoder_fn( + self._run_compilation, ) + self._precompile_input_embeddings_merger() + self._flush_compilations() + self._precompile_backbone_with_inputs_embeds() + self._flush_compilations() + if self.runner.scheduler_config.async_scheduling: + self._precompile_substitute_placeholder_token() + self._flush_compilations() + if self.runner.speculative_config: + self._precompile_subtract_num_rejected_tokens() + self._flush_compilations() + self._precompile_concat_last_sampled_tokens_and_draft_tokens( + ) + self._flush_compilations() + + if not self.runner.is_last_rank: + return + self._precompile_select_from_array() + self._flush_compilations() + if not self.runner.is_pooling_model: + self._precompile_compute_logits() + else: + self._precompile_compute_pooling() + self._flush_compilations() + # Skip sampling if already precompiled before KV cache allocation + if not self._sampling_precompiled: + self._precompile_sampling() + self._flush_compilations() + self._precompile_disagg_utils() + self._flush_compilations() + # Skip gather_logprobs if already precompiled before KV cache allocation + if not self._gather_logprobs_precompiled: + self._precompile_gather_logprobs() + self._flush_compilations() + self._precompile_structured_decoding() + self._flush_compilations() if self.runner.speculative_config: - self._precompile_subtract_num_rejected_tokens() - self._precompile_concat_last_sampled_tokens_and_draft_tokens( - ) - if not self.runner.is_last_rank: - return - self._precompile_select_from_array() - if not self.runner.is_pooling_model: - self._precompile_compute_logits() - else: - self._precompile_compute_pooling() - # Skip sampling if already precompiled before KV cache allocation - if not self._sampling_precompiled: - self._precompile_sampling() - self._precompile_disagg_utils() - # Skip gather_logprobs if already precompiled before KV cache allocation - if not self._gather_logprobs_precompiled: - self._precompile_gather_logprobs() - self._precompile_structured_decoding() - if self.runner.speculative_config: - self._precompile_speculative_decoding() - + self._precompile_speculative_decoding() + self._flush_compilations() + finally: + self._finalize_compilation() elapsed = time.perf_counter() - compilation_start_time self.runner.vllm_config.compilation_config.compilation_time += elapsed + def _finalize_compilation(self) -> None: + """Shut down the precompile pool and restore the thread stack default + so the bumped stack size doesn't leak to threads spawned later by the + engine.""" + if self._compile_executor is not None: + self._compile_executor.shutdown(wait=True) + self._compile_executor = None + if self._prev_stack_size is not None: + try: + threading.stack_size(self._prev_stack_size) + except (RuntimeError, ValueError): + pass + self._prev_stack_size = None + def _precompile_input_embeddings_merger(self) -> None: for num_tokens in self.runner.num_tokens_paddings: hidden_size = self.runner.vllm_config.model_config.get_hidden_size( @@ -286,26 +417,22 @@ def build_attn(block_tables: jax.Array | None) -> AttentionMetadata: for name in kv_cache_group.layer_names } - def model_fn_wrapper( - state_leaves, - kv_caches, - input_ids, - attention_metadata, - positions, - inputs_embeds, - layer_name_to_kvcache_index, - lora_metadata, - intermediate_tensors, - is_first_rank, - is_last_rank, - ): - kv_caches, hidden_states, *_ = self.runner.model_fn( - state_leaves, kv_caches, input_ids, attention_metadata, - inputs_embeds, positions, layer_name_to_kvcache_index, - lora_metadata, intermediate_tensors, is_first_rank, - is_last_rank) - self.runner.kv_caches = kv_caches - return hidden_states + def model_fn_warmup(_fn, _args, _call_kwargs): + out = self.runner.model_fn( + self.runner.state_leaves, + self.runner.kv_caches, + input_ids, + attention_metadata, + inputs_embeds, + positions, + tuple(self.runner.layer_name_to_kvcache_index.items()), + lora_metadata, + intermediate_tensors, + is_first_rank, + is_last_rank, + ) + self.runner.kv_caches = out[0] + return out with self.runner.maybe_select_dummy_loras( self.runner.lora_config, np.array([num_tokens], @@ -313,13 +440,13 @@ def model_fn_wrapper( lora_metadata = self.runner.lora_utils.extract_lora_metadata() self._run_compilation( name, - model_fn_wrapper, + self.runner.model_fn, self.runner.state_leaves, self.runner.kv_caches, input_ids, attention_metadata, - positions, inputs_embeds, + positions, tuple(self.runner.layer_name_to_kvcache_index.items()), lora_metadata, intermediate_tensors, @@ -327,6 +454,7 @@ def model_fn_wrapper( is_last_rank, num_tokens=num_tokens, num_reqs=num_reqs, + warmup_handler=model_fn_warmup, ) def _precompile_substitute_placeholder_token(self) -> None: @@ -621,11 +749,22 @@ def _precompile_select_from_array_helper( self._run_compilation( f"select_from_array [{name}]", - self.runner._select_from_array_fn, input_tensor, - indices_to_select, **{ + self.runner._select_from_array_fn, + self.runner, + input_tensor, + indices_to_select, + **{ "array_size": array_size, "index_size": indices_count - }) + }, + warmup_handler=self._skip_self_arg_warmup_handler) + + def _skip_self_arg_warmup_handler(self, fn, args, call_kwargs): + """Warmup handler for methods compiled with an explicit `self` as the + first positional arg. At warm-up time the object is already bound, so + we drop args[0] and forward the rest. + """ + return fn(*args[1:], **call_kwargs) def _precompile_select_from_array(self) -> None: logger.info("Compiling select_from_array with different input shapes.") @@ -769,6 +908,7 @@ def _precompile_sampling(self) -> None: sampling_metadata, num_reqs=num_reqs, do_sampling=do_sampling, + logprobs=logprobs, ) self._sampling_precompiled = True @@ -1210,20 +1350,10 @@ def _precompile_eagle3_helpers(self) -> None: padded_num_reqs=num_reqs, ) - def drafter_propose_fn_wrapper( - kv_caches, - input_ids, - attn_metadata, - last_token_indices, - target_hidden_states, - ): + def drafter_propose_warmup(_fn, _args, _call_kwargs): + new_args = (self.runner.kv_caches, ) + _args[1:] kv_caches, draft_token_ids = self.runner.drafter.propose( - kv_caches, - input_ids, - attn_metadata, - last_token_indices, - target_hidden_states, - ) + *new_args, **_call_kwargs) self.runner.kv_caches = kv_caches return draft_token_ids @@ -1237,13 +1367,14 @@ def drafter_propose_fn_wrapper( jnp.int32, dp_sharding) self._run_compilation( "drafter_propose", - drafter_propose_fn_wrapper, + self.runner.drafter.propose, self.runner.kv_caches, input_ids, attention_metadata, last_token_indices, draft_hidden_states, num_tokens=num_tokens, + warmup_handler=drafter_propose_warmup, ) aux_hidden_states = [ self._create_dummy_tensor( @@ -1363,20 +1494,10 @@ def _precompile_mtp_helpers(self) -> None: padded_num_reqs=num_reqs, ) - def drafter_propose_fn_wrapper( - kv_caches, - input_ids, - attn_metadata, - last_token_indices, - target_hidden_states, - ): + def drafter_propose_warmup(_fn, _args, _call_kwargs): + new_args = (self.runner.kv_caches, ) + _args[1:] kv_caches, draft_token_ids = self.runner.drafter.propose( - kv_caches, - input_ids, - attn_metadata, - last_token_indices, - target_hidden_states, - ) + *new_args, **_call_kwargs) self.runner.kv_caches = kv_caches return draft_token_ids @@ -1390,13 +1511,14 @@ def drafter_propose_fn_wrapper( jnp.int32, dp_sharding) self._run_compilation( "drafter_propose", - drafter_propose_fn_wrapper, + self.runner.drafter.propose, self.runner.kv_caches, input_ids, attention_metadata, last_token_indices, draft_hidden_states, num_tokens=num_tokens, + warmup_handler=drafter_propose_warmup, ) aux_hidden_states = (self._create_dummy_tensor( @@ -1455,9 +1577,11 @@ def _precompile_structured_decoding(self) -> None: self._run_compilation( "structured_decode", self.runner.structured_decoding_manager.structured_decode_fn, + self.runner.structured_decoding_manager, dummy_require_struct_decoding, dummy_grammar_bitmask, dummy_logits, arange, num_reqs=num_reqs, + warmup_handler=self._skip_self_arg_warmup_handler, ) diff --git a/tpu_inference/runner/tpu_runner.py b/tpu_inference/runner/tpu_runner.py index 2bda0331da..dc931d4012 100644 --- a/tpu_inference/runner/tpu_runner.py +++ b/tpu_inference/runner/tpu_runner.py @@ -982,23 +982,24 @@ def sample_tokens( ) self.execute_model_state = None - if grammar_output is not None: - ( - require_struct_decoding, grammar_bitmask_padded, arange - ) = self.structured_decoding_manager.prepare_structured_decoding_input( - logits, grammar_output) - logits = self.structured_decoding_manager.structured_decode_fn( - require_struct_decoding, - grammar_bitmask_padded, - logits, - arange, - ) - return self._sample_from_logits( - scheduler_output, attn_metadata, sampling_metadata, input_ids, - hidden_states, logits, aux_hidden_states, spec_decode_metadata, - kv_connector_output, logits_indices_selector, padded_num_reqs, - expert_indices, full_hidden_states, full_logits, req_ids_dp, - padded_num_scheduled_tokens_per_dp_rank) + with jax.set_mesh(self.mesh): + if grammar_output is not None: + ( + require_struct_decoding, grammar_bitmask_padded, arange + ) = self.structured_decoding_manager.prepare_structured_decoding_input( + logits, grammar_output) + logits = self.structured_decoding_manager.structured_decode_fn( + require_struct_decoding, + grammar_bitmask_padded, + logits, + arange, + ) + return self._sample_from_logits( + scheduler_output, attn_metadata, sampling_metadata, input_ids, + hidden_states, logits, aux_hidden_states, spec_decode_metadata, + kv_connector_output, logits_indices_selector, padded_num_reqs, + expert_indices, full_hidden_states, full_logits, req_ids_dp, + padded_num_scheduled_tokens_per_dp_rank) def _modify_prev_results(self): # If copy to host has not been done, we just wait. diff --git a/tpu_inference/worker/tpu_worker.py b/tpu_inference/worker/tpu_worker.py index a0781c7cd8..d4b20b703a 100644 --- a/tpu_inference/worker/tpu_worker.py +++ b/tpu_inference/worker/tpu_worker.py @@ -573,6 +573,7 @@ def initialize_from_config( and self.model_runner.model_config.enforce_eager)): self.model_runner.compilation_manager._precompile_sampling() self.model_runner.compilation_manager._precompile_gather_logprobs() + self.model_runner.compilation_manager._flush_compilations() # Init kv cache connector here, because it requires `kv_cache_config`. ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config) From 0d2c8c8f7c66818b876020d34bb93846187e31be Mon Sep 17 00:00:00 2001 From: patrickji2014 <110961369+patrickji2014@users.noreply.github.com> Date: Fri, 5 Jun 2026 17:44:48 -0700 Subject: [PATCH 18/37] Add owners for tools folder (#2832) Signed-off-by: patrickji2014 --- .github/CODEOWNERS | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 0e566625a8..360ef77976 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -63,7 +63,7 @@ /docker/ @jrplatin @QiliangCui @yiw-wang # Tools (kernel tuner, etc.) -/tools/ @vipannalla @kyuyeunk @patrickji2014 +/tools/ @vipannalla @kyuyeunk @patrickji2014 @yiw-wang @qiliangcui # Unowned: scripts, examples, and bot-updated support_matrices — # anyone can approve, no required CODEOWNER review. Bot pushes to From 39ba549082dcf758f2a12862953117626916980a Mon Sep 17 00:00:00 2001 From: Ming Huang Date: Fri, 5 Jun 2026 19:38:53 -0700 Subject: [PATCH 19/37] Fix data-parallel-size configuration typo for Gemma4 on v7x-2 (#2835) Signed-off-by: Ming Huang --- ...-dataset_custom-inlen_1024-outlen_500.json | 20 ++++++++++--------- ...-dataset_custom-inlen_1024-outlen_500.json | 10 ++++++---- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/.buildkite/benchmark/cases/daily/Gemma4-26B-dataset_custom-inlen_1024-outlen_500.json b/.buildkite/benchmark/cases/daily/Gemma4-26B-dataset_custom-inlen_1024-outlen_500.json index be6a8088e9..31cbdfbf35 100644 --- a/.buildkite/benchmark/cases/daily/Gemma4-26B-dataset_custom-inlen_1024-outlen_500.json +++ b/.buildkite/benchmark/cases/daily/Gemma4-26B-dataset_custom-inlen_1024-outlen_500.json @@ -24,16 +24,17 @@ "VLLM_USE_V1": "1", "MODEL_IMPL_TYPE": "flax_nnx", "USE_BATCHED_RPA_KERNEL": "1", - "MOE_REQUANTIZE_WEIGHT_DTYPE": "float8_e4m3fn" + "MOE_REQUANTIZE_WEIGHT_DTYPE": "float8_e4m3fn", + "VLLM_ENGINE_READY_TIMEOUT_S": "1800" }, "args": { "model": "google/gemma-4-26B-A4B-it", "max-num-seqs": 256, "max-num-batched-tokens": 4096, "data-parallel-size": { - "v6e-8": 2, - "v7x-2": 4, - "default": 1 + "v6e-8": 2, + "v7x-2": 1, + "default": 1 }, "tensor-parallel-size": { "v7x-2": 2, @@ -79,16 +80,17 @@ "VLLM_USE_V1": "1", "MODEL_IMPL_TYPE": "flax_nnx", "USE_BATCHED_RPA_KERNEL": "1", - "MOE_REQUANTIZE_WEIGHT_DTYPE": "float8_e4m3fn" + "MOE_REQUANTIZE_WEIGHT_DTYPE": "float8_e4m3fn", + "VLLM_ENGINE_READY_TIMEOUT_S": "1800" }, "args": { "model": "google/gemma-4-26B-A4B-it", "max-num-seqs": 256, "max-num-batched-tokens": 4096, "data-parallel-size": { - "v6e-8": 2, - "v7x-2": 4, - "default": 1 + "v6e-8": 2, + "v7x-2": 1, + "default": 1 }, "tensor-parallel-size": { "v7x-2": 2, @@ -121,4 +123,4 @@ } } ] -} +} \ No newline at end of file diff --git a/.buildkite/benchmark/cases/daily/Gemma4-31B-dataset_custom-inlen_1024-outlen_500.json b/.buildkite/benchmark/cases/daily/Gemma4-31B-dataset_custom-inlen_1024-outlen_500.json index 7834ff5301..d30e93f047 100644 --- a/.buildkite/benchmark/cases/daily/Gemma4-31B-dataset_custom-inlen_1024-outlen_500.json +++ b/.buildkite/benchmark/cases/daily/Gemma4-31B-dataset_custom-inlen_1024-outlen_500.json @@ -24,7 +24,8 @@ "env": { "VLLM_USE_V1": "1", "MODEL_IMPL_TYPE": "flax_nnx", - "USE_BATCHED_RPA_KERNEL": "1" + "USE_BATCHED_RPA_KERNEL": "1", + "VLLM_ENGINE_READY_TIMEOUT_S": "1800" }, "args": { "model": "google/gemma-4-31B-it", @@ -32,7 +33,7 @@ "max-num-batched-tokens": 4096, "data-parallel-size": { "v6e-8": 2, - "v7x-2": 4, + "v7x-2": 1, "default": 1 }, "tensor-parallel-size": { @@ -80,7 +81,8 @@ "env": { "VLLM_USE_V1": "1", "MODEL_IMPL_TYPE": "flax_nnx", - "USE_BATCHED_RPA_KERNEL": "1" + "USE_BATCHED_RPA_KERNEL": "1", + "VLLM_ENGINE_READY_TIMEOUT_S": "1800" }, "args": { "model": "google/gemma-4-31B-it", @@ -88,7 +90,7 @@ "max-num-batched-tokens": 4096, "data-parallel-size": { "v6e-8": 2, - "v7x-2": 4, + "v7x-2": 1, "default": 1 }, "tensor-parallel-size": { From 72b66c8933e283f005a9b7864037db684ccc7022 Mon Sep 17 00:00:00 2001 From: Chengji Yao Date: Fri, 5 Jun 2026 20:11:30 -0700 Subject: [PATCH 20/37] Fix mamba/GDN recurrent-state slot lifecycle leaks (#2779) Signed-off-by: Chengji Yao --- tests/runner/test_input_batch.py | 38 ++++ tests/runner/test_persistent_batch_manager.py | 179 ++++++++++++++++++ tpu_inference/runner/input_batch.py | 108 ++++++++++- .../runner/persistent_batch_manager.py | 58 +++++- 4 files changed, 375 insertions(+), 8 deletions(-) diff --git a/tests/runner/test_input_batch.py b/tests/runner/test_input_batch.py index 64e9d2487f..5372dd8059 100644 --- a/tests/runner/test_input_batch.py +++ b/tests/runner/test_input_batch.py @@ -282,6 +282,31 @@ def test_mamba_state_indices_freed_on_remove(input_batch: InputBatch): assert int(input_batch.mamba_state_indices_cpu[0]) == slot_first +def test_mamba_state_slot_preserved_when_request_unscheduled( + input_batch: InputBatch): + """Temporarily unscheduled requests must keep their physical mamba slot. + + The recurrent state lives in the mamba cache slot, not in the moving + persistent-batch row. Freeing the slot for an unscheduled-but-not-finished + request lets another request overwrite that state before the original + request resumes. + """ + req = create_dummy_request("req-0") + input_batch.add_request(req) + slot_first = int(input_batch.mamba_state_indices_cpu[0]) + + input_batch.remove_request("req-0", free_mamba_slot=False) + + assert req.mamba_state_slot == slot_first + assert int(input_batch.mamba_state_indices_cpu[0]) == 0 + assert all(slot_first not in pool + for pool in input_batch._free_mamba_slots_per_rank) + + input_batch.add_request(req) + + assert int(input_batch.mamba_state_indices_cpu[0]) == slot_first + + def test_mamba_state_indices_follow_condense(input_batch: InputBatch): """When condense moves a request to a different persistent-batch slot, its mamba state id must follow it — otherwise the GDN op reads stale @@ -543,6 +568,19 @@ def test_dp_mamba_remove_returns_slot_to_correct_rank( dp_input_batch._free_mamba_slots_per_rank[rank]) == pool_before +def test_dp_mamba_release_ignores_rank_null_slots(dp_input_batch: InputBatch): + """Rank-local null slots must never enter the reusable mamba slot pools.""" + local_slots = dp_input_batch._mamba_local_slots + pools_before = [ + list(pool) for pool in dp_input_batch._free_mamba_slots_per_rank + ] + + for rank in range(DP_SIZE): + dp_input_batch.release_mamba_slot(rank * local_slots) + + assert dp_input_batch._free_mamba_slots_per_rank == pools_before + + def test_dp_mamba_slots_unique_across_ranks(dp_input_batch: InputBatch): """Slots allocated to different ranks must never overlap.""" all_slots = [] diff --git a/tests/runner/test_persistent_batch_manager.py b/tests/runner/test_persistent_batch_manager.py index 5e40bea322..6159427edd 100644 --- a/tests/runner/test_persistent_batch_manager.py +++ b/tests/runner/test_persistent_batch_manager.py @@ -16,11 +16,58 @@ from unittest.mock import MagicMock, patch import numpy as np +from vllm.sampling_params import SamplingParams +from tpu_inference.runner.input_batch import CachedRequestState, InputBatch from tpu_inference.runner.persistent_batch_manager import \ PersistentBatchManager +def _create_cached_request(req_id: str) -> CachedRequestState: + return CachedRequestState( + req_id=req_id, + prompt_token_ids=[1, 2, 3], + mm_features=[], + sampling_params=SamplingParams(temperature=0.0), + pooling_params=None, + block_ids=[[1]], + num_computed_tokens=0, + lora_request=None, + output_token_ids=[], + ) + + +def _make_scheduler_output(*, + scheduled_req_ids=(), + preempted_req_ids=None, + resumed_req_ids=(), + new_block_ids=None): + scheduler_output = MagicMock() + req_data = MagicMock() + req_data.req_ids = list(scheduled_req_ids) + req_data.resumed_req_ids = set(resumed_req_ids) + req_data.num_computed_tokens = [0] * len(req_data.req_ids) + req_data.new_token_ids = [[] for _ in req_data.req_ids] + if new_block_ids is None: + req_data.new_block_ids = [None for _ in req_data.req_ids] + else: + req_data.new_block_ids = new_block_ids + req_data.num_output_tokens = [0] * len(req_data.req_ids) + scheduler_output.scheduled_cached_reqs = req_data + scheduler_output.scheduled_new_reqs = [] + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_scheduled_tokens = { + req_id: 1 + for req_id in scheduled_req_ids + } + scheduler_output.total_num_scheduled_tokens = len(scheduled_req_ids) + scheduler_output.finished_req_ids = set() + scheduler_output.free_encoder_mm_hashes = [] + scheduler_output.preempted_req_ids = set(preempted_req_ids or ()) + scheduler_output.assigned_dp_rank = {} + return scheduler_output + + class MockInputBatch: """Lightweight mock InputBatch that tracks the state needed by _reorder_batch: req_ids, request_distribution, and swap_states.""" @@ -128,6 +175,138 @@ def test_mixed_batch_needs_swap(self): class TestPersistentBatchManager(unittest.TestCase): + def test_update_states_preserves_mamba_slot_for_unscheduled_request(self): + req = _create_cached_request("req-0") + requests = {req.req_id: req} + input_batch = InputBatch( + max_num_reqs=4, + max_model_len=16, + max_num_batched_tokens=16, + pin_memory=False, + vocab_size=128, + block_sizes=[16], + ) + input_batch.add_request(req) + slot = int(input_batch.mamba_state_indices_cpu[0]) + + manager = PersistentBatchManager(requests, + input_batch, + encoder_cache={}, + uses_mrope=False, + model_config=MagicMock(), + is_last_rank=True) + + manager.update_states(_make_scheduler_output(), None) + + self.assertEqual(req.mamba_state_slot, slot) + self.assertFalse( + any(slot in pool + for pool in input_batch._free_mamba_slots_per_rank)) + + def test_update_states_releases_mamba_slot_for_preempted_request(self): + req = _create_cached_request("req-0") + requests = {req.req_id: req} + input_batch = InputBatch( + max_num_reqs=4, + max_model_len=16, + max_num_batched_tokens=16, + pin_memory=False, + vocab_size=128, + block_sizes=[16], + ) + input_batch.add_request(req) + slot = int(input_batch.mamba_state_indices_cpu[0]) + + manager = PersistentBatchManager(requests, + input_batch, + encoder_cache={}, + uses_mrope=False, + model_config=MagicMock(), + is_last_rank=True) + + manager.update_states( + _make_scheduler_output(preempted_req_ids={req.req_id}), None) + + self.assertIsNone(req.mamba_state_slot) + self.assertTrue( + any(slot in pool + for pool in input_batch._free_mamba_slots_per_rank)) + + def test_update_states_releases_preserved_mamba_slot_when_preempted(self): + req = _create_cached_request("req-0") + requests = {req.req_id: req} + input_batch = InputBatch( + max_num_reqs=4, + max_model_len=16, + max_num_batched_tokens=16, + pin_memory=False, + vocab_size=128, + block_sizes=[16], + ) + input_batch.add_request(req) + slot = int(input_batch.mamba_state_indices_cpu[0]) + + manager = PersistentBatchManager(requests, + input_batch, + encoder_cache={}, + uses_mrope=False, + model_config=MagicMock(), + is_last_rank=True) + + manager.update_states(_make_scheduler_output(), None) + self.assertEqual(req.mamba_state_slot, slot) + self.assertEqual(input_batch.num_reqs, 0) + self.assertFalse( + any(slot in pool + for pool in input_batch._free_mamba_slots_per_rank)) + + manager.update_states( + _make_scheduler_output(preempted_req_ids={req.req_id}), None) + + self.assertIsNone(req.mamba_state_slot) + self.assertTrue( + any(slot in pool + for pool in input_batch._free_mamba_slots_per_rank)) + + def test_update_states_resets_preserved_mamba_slot_when_resumed(self): + req = _create_cached_request("req-0") + requests = {req.req_id: req} + input_batch = InputBatch( + max_num_reqs=4, + max_model_len=16, + max_num_batched_tokens=16, + pin_memory=False, + vocab_size=128, + block_sizes=[16], + ) + input_batch.add_request(req) + slot = int(input_batch.mamba_state_indices_cpu[0]) + + manager = PersistentBatchManager(requests, + input_batch, + encoder_cache={}, + uses_mrope=False, + model_config=MagicMock(), + is_last_rank=True) + + manager.update_states(_make_scheduler_output(), None) + self.assertEqual(req.mamba_state_slot, slot) + + with patch.object(input_batch, + "release_mamba_slot", + wraps=input_batch.release_mamba_slot) as release: + manager.update_states( + _make_scheduler_output( + scheduled_req_ids=[req.req_id], + resumed_req_ids={req.req_id}, + new_block_ids=[[[2]]], + ), None) + + release.assert_called_with(slot) + self.assertEqual(input_batch.num_reqs, 1) + self.assertEqual(req.mamba_state_slot, + int(input_batch.mamba_state_indices_cpu[0])) + def test_update_states_pp_non_last_rank(self): """ the current rank is not the last rank. diff --git a/tpu_inference/runner/input_batch.py b/tpu_inference/runner/input_batch.py index eeeab9c03e..55a9411f3e 100644 --- a/tpu_inference/runner/input_batch.py +++ b/tpu_inference/runner/input_batch.py @@ -34,6 +34,7 @@ class CachedRequestState(NewRequestData): # Tuple of (token_ids, logprobs, ranks) numpy arrays, each of shape # [num_prompt_tokens-1, ...]. Set to None when prefill completes. in_progress_prompt_logprobs_cpu: Optional[tuple] = None + mamba_state_slot: Optional[int] = None def __post_init__(self): self.num_prompt_tokens = len(self.prompt_token_ids) @@ -190,6 +191,89 @@ def init_mamba_pools(self, mamba_num_blocks: int) -> None: self._free_mamba_slots_per_rank.append( list(range(base + self._mamba_local_slots - 1, base, -1))) + def release_mamba_slot(self, slot: Optional[int]) -> None: + if slot is None: + return + slot = int(slot) + if slot == 0 or slot % self._mamba_local_slots == 0: + return + rank = slot // self._mamba_local_slots + pool = self._free_mamba_slots_per_rank[rank] + if slot not in pool: + pool.append(slot) + + def assert_mamba_state_invariants( + self, + requests: Optional[dict[str, CachedRequestState]] = None, + assigned_dp_rank: Optional[dict[str, int]] = None, + ) -> None: + active_slots: list[int] = [] + active_req_ids = self._req_ids[:self.num_reqs] + free_slots = { + int(slot) + for pool in self._free_mamba_slots_per_rank + for slot in pool + } + + if sum(len(pool) + for pool in self._free_mamba_slots_per_rank) != len(free_slots): + raise AssertionError("Duplicate mamba slots in free pools") + + for req_index, req_id in enumerate(active_req_ids): + if req_id is None: + raise AssertionError( + f"Active mamba batch has a hole at index {req_index}") + slot = int(self.mamba_state_indices_cpu[req_index]) + active_slots.append(slot) + if slot <= 0 or slot % self._mamba_local_slots == 0: + raise AssertionError( + f"Request {req_id} has invalid mamba slot {slot}") + if slot in free_slots: + raise AssertionError( + f"Active request {req_id} uses free mamba slot {slot}") + if assigned_dp_rank is not None: + expected_rank = assigned_dp_rank.get(req_id, 0) + slot_rank = slot // self._mamba_local_slots + if slot_rank != expected_rank: + raise AssertionError( + f"Request {req_id} on DP rank {expected_rank} has " + f"mamba slot {slot} from rank {slot_rank}") + if requests is not None: + req_state = requests.get(req_id) + if req_state is not None and req_state.mamba_state_slot != slot: + raise AssertionError( + f"Request {req_id} active slot {slot} does not match " + f"cached slot {req_state.mamba_state_slot}") + + if len(set(active_slots)) != len(active_slots): + from collections import Counter + duplicate_active = { + slot + for slot, count in Counter(active_slots).items() if count > 1 + } + raise AssertionError( + f"Duplicate active mamba slots: {sorted(duplicate_active)}") + + tail = self.mamba_state_indices_cpu[self.num_reqs:] + if tail.any(): + nonzero_tail = sorted(set(tail[tail != 0].tolist())) + raise AssertionError( + f"Non-zero mamba slots in padded tail: {nonzero_tail}") + + if requests is not None: + preserved_slot_list = [ + int(req.mamba_state_slot) for req in requests.values() + if req.mamba_state_slot is not None + ] + preserved_slots = set(preserved_slot_list) + if len(preserved_slot_list) != len(preserved_slots): + raise AssertionError("Duplicate preserved mamba slots") + overlap = preserved_slots & free_slots + if overlap: + raise AssertionError( + f"Preserved mamba slots also in free pool: " + f"{sorted(overlap)}") + @property def req_ids(self) -> list[str]: # None elements should only be present transiently @@ -268,8 +352,19 @@ def add_request( # Allocate a fresh mamba state slot for this request. The slot stays # with the request through the persistent batch's lifetime, even when # condense moves the request to a different `req_index`. - self.mamba_state_indices_cpu[req_index] = ( - self._free_mamba_slots_per_rank[dp_rank].pop()) + if request.mamba_state_slot is None: + request.mamba_state_slot = self._free_mamba_slots_per_rank[ + dp_rank].pop() + else: + slot_rank = int( + request.mamba_state_slot) // self._mamba_local_slots + assert slot_rank == dp_rank, ( + f"Preserved mamba slot {request.mamba_state_slot} belongs to " + f"DP rank {slot_rank}, not {dp_rank}") + pool = self._free_mamba_slots_per_rank[dp_rank] + if request.mamba_state_slot in pool: + pool.remove(request.mamba_state_slot) + self.mamba_state_indices_cpu[req_index] = request.mamba_state_slot # NOTE(woosuk): self.generators should not include the requests that # do not have their own generator. @@ -351,7 +446,10 @@ def collect_sampling(sampling_params: SamplingParams) -> None: # No LoRA self.request_lora_mapping[req_index] = 0 - def remove_request(self, req_id: str) -> Optional[int]: + def remove_request(self, + req_id: str, + *, + free_mamba_slot: bool = True) -> Optional[int]: """This method must always be followed by a call to condense().""" req_index = self.req_id_to_index.pop(req_id, None) @@ -363,8 +461,8 @@ def remove_request(self, req_id: str) -> Optional[int]: # contents in the kv cache are stale and will be zeroed by the # has_initial_state guard when the next request takes this slot id. slot = int(self.mamba_state_indices_cpu[req_index]) - rank = slot // self._mamba_local_slots - self._free_mamba_slots_per_rank[rank].append(slot) + if free_mamba_slot: + self.release_mamba_slot(slot) # Clear this position to slot 0 (the null block) so the trailing # tail of `mamba_state_indices_cpu` (which the GDN op reads over # its full length every step) cannot alias an active slot. diff --git a/tpu_inference/runner/persistent_batch_manager.py b/tpu_inference/runner/persistent_batch_manager.py index bf4d22bf5d..2ec067a758 100644 --- a/tpu_inference/runner/persistent_batch_manager.py +++ b/tpu_inference/runner/persistent_batch_manager.py @@ -98,8 +98,9 @@ def update_states(self, scheduler_output: "VllmSchedulerOutput", If False, we can skip copying SamplingMetadata to the TPU. """ # Remove finished requests from the cached states. + finished_req_states = {} for req_id in scheduler_output.finished_req_ids: - self.requests.pop(req_id, None) + finished_req_states[req_id] = self.requests.pop(req_id, None) # Remove the finished requests from the persistent batch. # NOTE(woosuk): There could be an edge case where finished_req_ids and @@ -111,7 +112,16 @@ def update_states(self, scheduler_output: "VllmSchedulerOutput", for req_id in scheduler_output.finished_req_ids: req_index = self.input_batch.remove_request(req_id) if req_index is not None: + req_state = finished_req_states.get(req_id) + if req_state is not None: + req_state.mamba_state_slot = None removed_req_indices.append(req_index) + else: + req_state = finished_req_states.get(req_id) + if req_state is not None: + self.input_batch.release_mamba_slot( + req_state.mamba_state_slot) + req_state.mamba_state_slot = None # Free the cached encoder outputs. for mm_hash in scheduler_output.free_encoder_mm_hashes: @@ -124,13 +134,50 @@ def update_states(self, scheduler_output: "VllmSchedulerOutput", # they will be scheduled again sometime in the future. scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys() cached_req_ids = self.input_batch.req_id_to_index.keys() - unscheduled_req_ids = cached_req_ids - scheduled_req_ids + resumed_req_ids = set( + getattr(scheduler_output.scheduled_cached_reqs, "resumed_req_ids", + ()) or ()) + preempted_req_ids = set( + getattr(scheduler_output, "preempted_req_ids", ()) or ()) + reset_mamba_req_ids = preempted_req_ids | resumed_req_ids + + # A request can be temporarily removed from the persistent batch while + # keeping its physical mamba slot. If it is then preempted or resumed + # before being re-added, the active-batch removal loop below cannot see + # it. Reset the preserved slot here so a recomputed request cannot read + # stale recurrent state from its old slot. + for req_id in reset_mamba_req_ids: + if req_id in self.input_batch.req_id_to_index: + continue + req_state = self.requests.get(req_id) + if req_state is None: + continue + self.input_batch.release_mamba_slot(req_state.mamba_state_slot) + req_state.mamba_state_slot = None + + # Usually resumed requests are not present in the persistent batch. + # Forced preemption can make a resumed request still appear there; clear + # it first so it is re-added through the normal resumed-request path. + unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - + resumed_req_ids) # NOTE(woosuk): The persistent batch optimization assumes that # consecutive batches contain mostly the same requests. If batches # have low request overlap (e.g., alternating between two distinct # sets of requests), this optimization becomes very inefficient. for req_id in unscheduled_req_ids: - req_index = self.input_batch.remove_request(req_id) + req_index = self.input_batch.req_id_to_index.get(req_id) + reset_mamba_slot = (req_id in preempted_req_ids + or req_id in resumed_req_ids) + if req_index is not None: + req_state = self.requests.get(req_id) + if reset_mamba_slot: + if req_state is not None: + req_state.mamba_state_slot = None + else: + self.requests[req_id].mamba_state_slot = int( + self.input_batch.mamba_state_indices_cpu[req_index]) + req_index = self.input_batch.remove_request( + req_id, free_mamba_slot=reset_mamba_slot) assert req_index is not None removed_req_indices.append(req_index) @@ -253,6 +300,8 @@ def update_states(self, scheduler_output: "VllmSchedulerOutput", # The smaller empty indices are filled first. removed_req_indices = sorted(removed_req_indices, reverse=True) dp_rank_map = getattr(scheduler_output, 'assigned_dp_rank', None) + if not isinstance(dp_rank_map, dict): + dp_rank_map = None for req_id in req_ids_to_add: req_state = self.requests[req_id] if removed_req_indices: @@ -271,4 +320,7 @@ def update_states(self, scheduler_output: "VllmSchedulerOutput", batch_changed = len(unscheduled_req_ids) > 0 or len(req_ids_to_add) > 0 # TODO(jevinjiang): I assume we do not need to set batch_changed to true if just swapping requests. self._reorder_batch(scheduler_output) + if isinstance(self.input_batch, InputBatch): + self.input_batch.assert_mamba_state_invariants( + self.requests, dp_rank_map) return batch_changed From 9006eb65e973f9ba7a8ded3a62e2b7894833ddce Mon Sep 17 00:00:00 2001 From: patrickji2014 <110961369+patrickji2014@users.noreply.github.com> Date: Fri, 5 Jun 2026 20:32:52 -0700 Subject: [PATCH 21/37] Make Kernel Tuner Result Inspector CLI Support Filter and Print as Table Format (#2675) Signed-off-by: patrick ji Signed-off-by: patrickji2014 --- .buildkite/pipeline_jax.yml | 2 +- tools/kernel/tuner/v1/README.md | 35 +- tools/kernel/tuner/v1/inspect_result_cli.py | 468 +++++++++++++++++- .../tuner/v1/tests/test_inspect_result_cli.py | 143 ++++++ 4 files changed, 622 insertions(+), 26 deletions(-) create mode 100644 tools/kernel/tuner/v1/tests/test_inspect_result_cli.py diff --git a/.buildkite/pipeline_jax.yml b/.buildkite/pipeline_jax.yml index 66ae55d897..e9d1ff6d3d 100644 --- a/.buildkite/pipeline_jax.yml +++ b/.buildkite/pipeline_jax.yml @@ -679,7 +679,7 @@ steps: - | if [[ "$TPU_VERSION" == "tpu7x" ]]; then .buildkite/scripts/run_in_docker.sh bash -c " - TPU_VERSION=\${TPU_VERSION} TPU_CORES=\${TPU_CORES} python -m pytest tools/kernel/tuner/v1/tests/test_kernel_tuner_runner.py -v -rs" + TPU_VERSION=\${TPU_VERSION} TPU_CORES=\${TPU_CORES} python -m pytest tools/kernel/tuner/v1/tests -v -rs" else echo "Skipping: Step only needs to run on TPU v7x." exit 0 diff --git a/tools/kernel/tuner/v1/README.md b/tools/kernel/tuner/v1/README.md index 11cc290a46..09e08eacd8 100644 --- a/tools/kernel/tuner/v1/README.md +++ b/tools/kernel/tuner/v1/README.md @@ -327,17 +327,38 @@ inspect|cs=testing_tuning_infra_11|run=001> query_run_status #### Query minimum latency results ``` -query_min_latency [--case_set_id ID] [--run_id ID] +query_min_latency [--case_set_id ID] [--run_id ID] [--show FIELD ...] ``` -For each unique `TuningKey`, shows the best measured latency and the corresponding `TunableParam` configuration. +For each unique `TuningKey`, shows the best measured latency and the corresponding `TunableParam` configuration. If repeatable --show option is specified, only the FIELDs are shown. Without --show option, all the fields in TuningKey and TunableParams are shown as a table. ``` -inspect|cs=testing_tuning_infra_11|run=001> query_min_latency - tuning_key={"key1": 1, "key2": 4} best_latency_us=135516 warmup_us=189835 tunable_params={"param1": 7, "param2": 11} case_id=1 - tuning_key={"key1": 1, "key2": 5} best_latency_us=81849 warmup_us=108349 tunable_params={"param1": 7, "param2": 10} case_id=2 - tuning_key={"key1": 2, "key2": 4} best_latency_us=79251 warmup_us=180675 tunable_params={"param1": 7, "param2": 10} case_id=4 - tuning_key={"key1": 2, "key2": 5} best_latency_us=105251 warmup_us=49820 tunable_params={"param1": 7, "param2": 11} case_id=7 +inspect|cs=mla_tuning_0|run=4> query_min_latency --show max_num_tokens --show actual_num_q_heads --show actual_lkv_dim --show actual_r_dim --show decode_batch_size --show num_kv_pages_per_block --show latency_us +max_num_tokens actual_num_q_heads actual_lkv_dim actual_r_dim decode_batch_size num_kv_pages_per_block latency_us +-------------- ------------------ -------------- ------------ ----------------- ---------------------- ---------- +128 128 512 64 16 1 2059 +... +64 128 512 64 16 1 2041 +8 128 512 64 8 2 2035 +``` + +#### Query case latency + +``` +query_case_latency Query latency for tuning cases with optional field filters + (--case_set_id ID --run_id ID [--filter_key FIELD=VALUE ...] [--show FIELD ...] [--show_all]) +``` + +FIELD can be any key in tuning_key or tunable_params. --show option behaves the same as above. --show_all includes all cases, even ones where tuning failed. + +``` +inspect|cs=mla_tuning_0|run=4> query_case_latency --filter_key max_num_tokens=4 --show max_num_tokens --show actual_num_q_heads --show actual_lkv_dim --show actual_r_dim --show decode_batch_size --show num_kv_pages_per_block --show latency_us --show_all +max_num_tokens actual_num_q_heads actual_lkv_dim actual_r_dim decode_batch_size num_kv_pages_per_block latency_us +-------------- ------------------ -------------- ------------ ----------------- ---------------------- ---------- +4 128 512 64 16 1 2078 +4 128 512 64 8 1 2111 +... +4 128 512 64 32 1 FAILURE ``` #### Other diff --git a/tools/kernel/tuner/v1/inspect_result_cli.py b/tools/kernel/tuner/v1/inspect_result_cli.py index 9b118d1429..b26ebe8b01 100644 --- a/tools/kernel/tuner/v1/inspect_result_cli.py +++ b/tools/kernel/tuner/v1/inspect_result_cli.py @@ -22,13 +22,21 @@ count_buckets Count buckets (--case_set_id ID --run_id ID) list_bucket_status Show completed vs pending counts (--case_set_id ID --run_id ID) query_run_status Show timing info for a run (--case_set_id ID --run_id ID) - query_min_latency Show best latency per TuningKey (--case_set_id ID --run_id ID) + query_min_latency Show best latency per TuningKey (--case_set_id ID --run_id ID [--show FIELD ...]) + query_case_latency Query latency for tuning cases with optional field filters + (--case_set_id ID --run_id ID [--filter_key FIELD=VALUE ...] [--show FIELD ...] [--show_all]) + FIELD can be any key in tuning_key or tunable_params. show_all includes unsuccessful cases; + By default only successful cases are shown. """ import argparse +import ast +import atexit import json +import math import os from collections import defaultdict +from enum import Enum # --------------------------------------------------------------------------- # Local backend helpers @@ -178,6 +186,149 @@ def local_query_min_latency(db_path, case_set_id, run_id): key=lambda x: json.dumps(x['tuning_key'], sort_keys=True)) +# --------------------------------------------------------------------------- +# Filtering helpers +# --------------------------------------------------------------------------- +class FilterResult(Enum): + MATCH = 1 + NO_MATCH = 2 + INVALID_FILTER = 3 + + +def _matches_filter(kv: dict, filter_keys: list) -> FilterResult: + """Return FilterResult.MATCH if a CaseKeyValue dict passes all KEY=VALUE filters. Return + FilterResult.NO_MATCH if any filter does not match, or FilterResult.INVALID_FILTER if any + filter is malformed. + + filter_keys is a list of strings like ["max_num_tokens=4", "q_dtype=fp8"]. + Fields are looked up in both tuning_key and tunable_params sub-dicts. + Type coercion is attempted to match the stored value's type. + """ + combined = {} + combined.update(kv.get('tuning_key') or {}) + combined.update(kv.get('tunable_params') or {}) + + for kv_str in filter_keys: + if '=' not in kv_str: + print( + f'Warning: invalid filter "{kv_str}" ignored (expected format FIELD=VALUE)' + ) + return FilterResult.INVALID_FILTER + field, raw = kv_str.split('=', 1) + field = field.strip() + raw = raw.strip() + if field not in combined: + print( + f'Warning: Filter field "{field}" not found in tuning_key or tunable_params' + ) + return FilterResult.INVALID_FILTER + stored = combined[field] + if stored is None: + if raw.lower() not in ('none', 'null', ''): + return FilterResult.NO_MATCH + elif isinstance(stored, bool): + if raw.lower() not in ('true', 'false', '1', '0', 'yes', 'no'): + print( + f'Warning: Invalid boolean value "{raw}" for field "{field}"' + ) + return FilterResult.INVALID_FILTER + if stored != (raw.lower() in ('true', '1', 'yes')): + return FilterResult.NO_MATCH + elif isinstance(stored, int): + try: + if stored != int(raw): + return FilterResult.NO_MATCH + except ValueError: + return FilterResult.INVALID_FILTER + elif isinstance(stored, float): + try: + if not math.isclose(stored, float(raw), rel_tol=1e-9): + return FilterResult.NO_MATCH + except ValueError: + return FilterResult.INVALID_FILTER + elif isinstance(stored, list): + try: + if stored != list(ast.literal_eval(raw)): + return FilterResult.NO_MATCH + except Exception: # pylint: disable=broad-except + return FilterResult.INVALID_FILTER + else: + if str(stored) != raw: + return FilterResult.NO_MATCH + return FilterResult.MATCH + + +def row_sort_key(row): + status = row.get('ProcessedStatus') + lat = row.get('Latency') + if status == 'SUCCESS': + return (0, lat, row.get('CaseId')) + return (1, float('inf'), row.get('CaseId')) + + +def local_query_case_latency(db_path, + case_set_id, + run_id, + filter_keys=None, + show_all=False): + """Return case latency rows matching filters. + + By default, only successful rows are returned. + If show_all is True, unsuccessful rows are also included. + """ + results = _read_json(db_path, 'CaseResults') + cases = _read_json(db_path, 'KernelTuningCases') + + case_kv_map = { + (c['ID'], c['CaseId']): c.get('CaseKeyValue') + for c in cases + } + + relevant = [ + r for r in results + if r['ID'] == case_set_id and str(r['RunId']) == str(run_id) + ] + if not show_all: + relevant = [ + r for r in relevant + if r.get('ProcessedStatus') == 'SUCCESS' and r.get('Latency') + ] + + rows = [] + for r in relevant: + kv_str = case_kv_map.get((r['ID'], r['CaseId'])) + if not kv_str: + print( + f'Warning: no CaseKeyValue found for CaseId={r["CaseId"]}; skipping' + ) + continue + try: + kv = json.loads(kv_str) + except (json.JSONDecodeError, TypeError): + print( + f'Warning: failed to decode CaseKeyValue for CaseId={r["CaseId"]}; skipping' + ) + continue + if filter_keys: + result = _matches_filter(kv, filter_keys) + if result == FilterResult.INVALID_FILTER: + print('One or more invalid filters; aborting query.') + return [] + if result != FilterResult.MATCH: + continue + rows.append({ + 'tuning_key': kv.get('tuning_key'), + 'tunable_params': kv.get('tunable_params'), + 'ProcessedStatus': r.get('ProcessedStatus'), + 'Latency': r.get('Latency'), + 'WarmupTime': r.get('WarmupTime'), + 'TotalTime': r.get('TotalTime'), + 'CaseId': r.get('CaseId'), + }) + + return sorted(rows, key=row_sort_key) + + # --------------------------------------------------------------------------- # Spanner backend helpers # --------------------------------------------------------------------------- @@ -356,6 +507,63 @@ def spanner_query_min_latency(db, case_set_id, run_id): key=lambda x: json.dumps(x['tuning_key'], sort_keys=True)) +def spanner_query_case_latency(db, + case_set_id, + run_id, + filter_keys=None, + show_all=False): + """Return case latency rows matching filters. + + By default, only successful rows are returned. + If show_all is True, unsuccessful rows are also included. + """ + from google.cloud import \ + spanner as gspanner # pylint: disable=import-outside-toplevel + where_status = "" if show_all else "AND cr.ProcessedStatus = 'SUCCESS'" + query = f""" + SELECT cr.CaseId, cr.Latency, cr.WarmupTime, cr.TotalTime, + cr.ProcessedStatus, ktc.CaseKeyValue + FROM CaseResults cr + JOIN KernelTuningCases ktc ON cr.ID = ktc.ID AND cr.CaseId = ktc.CaseId + WHERE cr.ID = @id AND cr.RunId = @rid {where_status} + ORDER BY cr.Latency + """ + rows = [] + with db.snapshot() as snap: + for case_id, lat, warmup, total_time, status, kv_str in snap.execute_sql( + query, + params={ + 'id': case_set_id, + 'rid': run_id + }, + param_types={ + 'id': gspanner.param_types.STRING, + 'rid': gspanner.param_types.STRING, + }): + try: + kv = json.loads(kv_str) + except (json.JSONDecodeError, TypeError): + continue + if filter_keys: + result = _matches_filter(kv, filter_keys) + if result == FilterResult.INVALID_FILTER: + print('One or more invalid filters; aborting query.') + return [] + if result != FilterResult.MATCH: + continue + rows.append({ + 'tuning_key': kv.get('tuning_key'), + 'tunable_params': kv.get('tunable_params'), + 'ProcessedStatus': status, + 'Latency': lat, + 'WarmupTime': warmup, + 'TotalTime': total_time, + 'CaseId': case_id, + }) + + return sorted(rows, key=row_sort_key) + + # --------------------------------------------------------------------------- # Display # --------------------------------------------------------------------------- @@ -378,16 +586,110 @@ def _print_table(rows, headers=None): print(fmt.format(*[str(row.get(h, '')) for h in headers])) -def _print_min_latency(rows): +def _print_flattened_table(rows, + builtin_cols, + row_builder, + show_fields=None, + empty_msg=' (no results)', + count_suffix=''): + """Shared helper: flatten rows, resolve show_fields, and print a table. + + Args: + rows: raw result rows. + builtin_cols: ordered list of column names produced by row_builder. + row_builder: callable(row) -> dict containing exactly the builtin cols. + tuning_key and tunable_params are merged in automatically after. + show_fields: optional list of columns to display (subset of all_cols). + empty_msg: message printed when rows is empty. + count_suffix: appended to the trailing "(N result(s))" line. + """ if not rows: - print(' (no successful results)') + print(empty_msg) return + + flat_rows = [] + all_extra = [] + seen_extra = set() + colliding_fields = set() for r in rows: - print(f" tuning_key={json.dumps(r['tuning_key'])}" - f" best_latency_us={r['Latency']}" - f" warmup_us={r['WarmupTime']}" - f" tunable_params={json.dumps(r['tunable_params'])}" - f" case_id={r['CaseId']}") + flat = row_builder(r) + dynamic_fields = {} + dynamic_fields.update(r.get('tuning_key') or {}) + dynamic_fields.update(r.get('tunable_params') or {}) + + for field in list(dynamic_fields): + if field in builtin_cols: + colliding_fields.add(field) + del dynamic_fields[field] + continue + if field not in seen_extra: + seen_extra.add(field) + all_extra.append(field) + + flat.update(dynamic_fields) + flat_rows.append(flat) + + if colliding_fields: + print(' Warning: dynamic field(s) ignored due to built-in column ' + f'name collision: {", ".join(sorted(colliding_fields))}') + + all_cols = builtin_cols + all_extra + + if show_fields: + unknown = [f for f in show_fields if f not in all_cols] + if unknown: + print(f' Warning: unknown field(s) ignored: {", ".join(unknown)}') + cols = [f for f in show_fields if f in all_cols] + if not cols: + print(' (no valid --show fields; showing all columns)') + cols = all_cols + else: + cols = all_cols + + _print_table(flat_rows, headers=cols) + print(f' ({len(rows)} result(s){count_suffix})') + + +def _print_min_latency(rows, show_fields=None): + """Print query_min_latency results as a table.""" + _print_flattened_table( + rows, + builtin_cols=['case_id', 'latency_us', 'warmup_us'], + row_builder=lambda r: { + 'case_id': r['CaseId'], + 'latency_us': r['Latency'], + 'warmup_us': r['WarmupTime'], + }, + show_fields=show_fields, + empty_msg=' (no successful results)', + ) + + +def _print_case_latency(rows, show_fields=None): + """Print query_case_latency results as a table.""" + + def _build_row(r): + status = r.get('ProcessedStatus') + is_success = status == 'SUCCESS' + return { + 'case_id': r['CaseId'], + 'processed_status': status, + 'latency_us': r['Latency'] if is_success else 'FAILURE', + 'warmup_us': r['WarmupTime'] if is_success else 'FAILURE', + 'total_time_us': r.get('TotalTime'), + } + + _print_flattened_table( + rows, + builtin_cols=[ + 'case_id', 'processed_status', 'latency_us', 'warmup_us', + 'total_time_us' + ], + row_builder=_build_row, + show_fields=show_fields, + empty_msg=' (no matching results)', + count_suffix=', sorted by latency', + ) # --------------------------------------------------------------------------- @@ -450,6 +752,68 @@ def _build_parser(): help='Show best latency per TuningKey.') p.add_argument('--case_set_id', required=True) p.add_argument('--run_id', required=True) + p.add_argument( + '--show', + dest='show_fields', + action='append', + default=None, + metavar='FIELD', + help=('Only display this column in the output table. ' + 'Repeat to show multiple columns. ' + 'Built-in columns: case_id, latency_us, warmup_us. ' + 'Any tuning_key or tunable_params field name is also valid. ' + 'Example: --show latency_us --show max_num_tokens'), + ) + + p = sub.add_parser( + 'query_case_latency', + help='Query latency for tuning cases, with optional field filters.', + description= + ('Show latency for all successful tuning cases matching the given filters.\n' + 'Use --filter_key FIELD=VALUE (repeatable) to filter by any field of\n' + 'tuning_key or tunable_params in the stored CaseKeyValue.\n\n' + 'FIELD can be any key present in the tuning_key or tunable_params\n' + 'sub-dicts of the case — field names vary by case set type.\n\n' + 'Example:\n' + ' query_case_latency --case_set_id X --run_id Y \\\n' + ' --filter_key max_num_tokens=4 --filter_key q_dtype=fp8'), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + p.add_argument('--case_set_id', required=True) + p.add_argument('--run_id', required=True) + p.add_argument( + '--filter_key', + dest='filter_keys', + action='append', + default=[], + metavar='FIELD=VALUE', + help= + ('Filter by a TuningKey or TunableParams field. ' + 'Repeat for multiple filters. Example: --filter_key max_num_tokens=4' + ), + ) + + p.add_argument( + '--show', + dest='show_fields', + action='append', + default=None, + metavar='FIELD', + help= + ('Only display this column in the output table. ' + 'Repeat to show multiple columns. ' + 'Built-in columns: case_id, processed_status, latency_us, ' + 'warmup_us, total_time_us. ' + 'Any tuning_key or tunable_params field name is also valid. ' + 'Example: --show latency_us --show max_num_tokens --show decode_batch_size' + ), + ) + p.add_argument( + '--show_all', + action='store_true', + help=('Include unsuccessful case results as well. ' + 'By default, only successful results are shown.'), + ) return parser @@ -546,9 +910,19 @@ def _run_command(args, source, db_path=None, spanner_db=None): print(f' {k}: {v}') elif args.command == 'query_min_latency': - _print_min_latency( - local_query_min_latency(db_path, args.case_set_id, - args.run_id)) + _print_min_latency(local_query_min_latency(db_path, + args.case_set_id, + args.run_id), + show_fields=args.show_fields) + + elif args.command == 'query_case_latency': + _print_case_latency(local_query_case_latency( + db_path, + args.case_set_id, + args.run_id, + filter_keys=args.filter_keys, + show_all=args.show_all), + show_fields=args.show_fields) else: # spanner if args.command == 'list_case_sets': @@ -587,9 +961,18 @@ def _run_command(args, source, db_path=None, spanner_db=None): print(f' {k}: {v}') elif args.command == 'query_min_latency': - _print_min_latency( - spanner_query_min_latency(spanner_db, args.case_set_id, - args.run_id)) + _print_min_latency(spanner_query_min_latency( + spanner_db, args.case_set_id, args.run_id), + show_fields=args.show_fields) + + elif args.command == 'query_case_latency': + _print_case_latency(spanner_query_case_latency( + spanner_db, + args.case_set_id, + args.run_id, + filter_keys=args.filter_keys, + show_all=args.show_all), + show_fields=args.show_fields) # --------------------------------------------------------------------------- @@ -605,21 +988,70 @@ def _run_command(args, source, db_path=None, spanner_db=None): count_buckets [--case_set_id ID] [--run_id ID] list_bucket_status [--case_set_id ID] [--run_id ID] query_run_status [--case_set_id ID] [--run_id ID] - query_min_latency [--case_set_id ID] [--run_id ID] + query_min_latency [--case_set_id ID] [--run_id ID] [--show FIELD ...] + --show: columns to display (default: all). Built-ins: case_id, latency_us, warmup_us + query_case_latency [--case_set_id ID] [--run_id ID] [--filter_key FIELD=VALUE ...] [--show FIELD ...] [--show_all] + --filter_key: any key in tuning_key or tunable_params (varies by case set type) + --show: columns to display (default: all). Built-ins: case_id, processed_status, + latency_us, warmup_us, total_time_us + --show_all: include unsuccessful rows (default shows only successful rows) + Example: query_case_latency --show_all --show processed_status --show latency_us + Use Up/Down arrows to recall command history help exit / quit """ +def _setup_console_history(): + """Enable persistent command history for interactive console mode. + + Returns: + Path to the history log file if readline is available, else None. + """ + try: + import readline # pylint: disable=import-outside-toplevel + except ImportError: + print( + 'Warning: readline module not available; command history disabled.' + ) + return None + + # Allow overriding location, but default to a clear log-like filename. + history_path = os.path.expanduser( + os.environ.get('INSPECT_RESULT_CLI_HISTORY_FILE', + '~/.inspect_result_cli_history.log')) + try: + if os.path.exists(history_path): + readline.read_history_file(history_path) + except OSError: + print( + f'Warning: could not read history file at {history_path}; starting with empty history.' + ) + + readline.set_history_length(2000) + + def _save_history(): + try: + readline.write_history_file(history_path) + except OSError: + print(f'Warning: could not write history file at {history_path}.') + + atexit.register(_save_history) + return history_path + + def _console_loop(source, db_path, spanner_db, global_args): """Run an interactive REPL until the user types exit/quit.""" parser = _build_parser() session_case_set_id = None session_run_id = None + history_path = _setup_console_history() print('\nKernel Tuning Inspector — console mode') print(f'Source: {source}' + (f' DB: {db_path}' if source == 'local' else '')) + if history_path: + print(f'History log: {history_path}') print(_COMMANDS_HELP) def _prompt(): @@ -653,8 +1085,8 @@ def _prompt(): else: new_case_set_id = tokens[1] if new_case_set_id != session_case_set_id: - session_run_id = None if session_run_id is not None: + session_run_id = None print(' run_id cleared') session_case_set_id = new_case_set_id print(f' case_set_id set to: {session_case_set_id}') @@ -671,11 +1103,11 @@ def _prompt(): # that actually accept those flags. _cmds_with_case_set_id = { 'list_runs', 'count_buckets', 'list_bucket_status', - 'query_run_status', 'query_min_latency' + 'query_run_status', 'query_min_latency', 'query_case_latency' } _cmds_with_run_id = { 'count_buckets', 'list_bucket_status', 'query_run_status', - 'query_min_latency' + 'query_min_latency', 'query_case_latency' } cmd = tokens[0] if '--case_set_id' not in line and session_case_set_id is not None \ diff --git a/tools/kernel/tuner/v1/tests/test_inspect_result_cli.py b/tools/kernel/tuner/v1/tests/test_inspect_result_cli.py new file mode 100644 index 0000000000..8384028cea --- /dev/null +++ b/tools/kernel/tuner/v1/tests/test_inspect_result_cli.py @@ -0,0 +1,143 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +from tools.kernel.tuner.v1.inspect_result_cli import (FilterResult, + _matches_filter) + + +class TestMatchesFilter(unittest.TestCase): + + def setUp(self): + # A sample CaseKeyValue dictionary containing both tuning_key and tunable_params + self.sample_kv = { + "tuning_key": { + "max_num_tokens": 4, + "q_dtype": "fp8", + "use_bias": True, + "threshold": 0.5, + "empty_val": None + }, + "tunable_params": { + "block_sizes": [16, 32], + "name": "test_case", + "is_active": False + } + } + + # --- Match & No Match Tests --- + + def test_integer_filtering(self): + self.assertEqual(_matches_filter(self.sample_kv, ["max_num_tokens=4"]), + FilterResult.MATCH) + self.assertEqual(_matches_filter(self.sample_kv, ["max_num_tokens=8"]), + FilterResult.NO_MATCH) + + def test_string_filtering(self): + self.assertEqual(_matches_filter(self.sample_kv, ["q_dtype=fp8"]), + FilterResult.MATCH) + self.assertEqual(_matches_filter(self.sample_kv, ["q_dtype=bf16"]), + FilterResult.NO_MATCH) + + def test_boolean_filtering_true(self): + # The function accepts 'true', '1', and 'yes' as True equivalents + for val in ["true", "1", "yes", "TRUE"]: + self.assertEqual( + _matches_filter(self.sample_kv, [f"use_bias={val}"]), + FilterResult.MATCH) + self.assertEqual(_matches_filter(self.sample_kv, ["use_bias=false"]), + FilterResult.NO_MATCH) + + def test_boolean_filtering_false(self): + # The function accepts 'false', '0', and 'no' as False equivalents + for val in ["false", "0", "no", "FALSE"]: + self.assertEqual( + _matches_filter(self.sample_kv, [f"is_active={val}"]), + FilterResult.MATCH) + self.assertEqual(_matches_filter(self.sample_kv, ["is_active=true"]), + FilterResult.NO_MATCH) + + def test_float_filtering(self): + self.assertEqual(_matches_filter(self.sample_kv, ["threshold=0.5"]), + FilterResult.MATCH) + self.assertEqual(_matches_filter(self.sample_kv, ["threshold=0.6"]), + FilterResult.NO_MATCH) + + def test_list_filtering(self): + self.assertEqual( + _matches_filter(self.sample_kv, ["block_sizes=[16, 32]"]), + FilterResult.MATCH) + # Should be NO_MATCH if the list doesn't match exactly + self.assertEqual(_matches_filter(self.sample_kv, ["block_sizes=[16]"]), + FilterResult.NO_MATCH) + + def test_none_filtering(self): + # The function accepts 'none', 'null', or empty string as None equivalents + for val in ["none", "null", "", "NONE"]: + self.assertEqual( + _matches_filter(self.sample_kv, [f"empty_val={val}"]), + FilterResult.MATCH) + self.assertEqual(_matches_filter(self.sample_kv, ["empty_val=0"]), + FilterResult.NO_MATCH) + + def test_multiple_filters(self): + # Match if all are true + self.assertEqual( + _matches_filter(self.sample_kv, + ["max_num_tokens=4", "q_dtype=fp8"]), + FilterResult.MATCH) + # No match if even one is false + self.assertEqual( + _matches_filter(self.sample_kv, + ["max_num_tokens=4", "q_dtype=bf16"]), + FilterResult.NO_MATCH) + + def test_empty_filter_list(self): + # Should default to MATCH if no filters are applied + self.assertEqual(_matches_filter(self.sample_kv, []), + FilterResult.MATCH) + + # --- Invalid Filter Tests --- + + def test_missing_equals_sign(self): + self.assertEqual(_matches_filter(self.sample_kv, ["max_num_tokens4"]), + FilterResult.INVALID_FILTER) + + def test_missing_field(self): + self.assertEqual( + _matches_filter(self.sample_kv, ["non_existent_field=1"]), + FilterResult.INVALID_FILTER) + + def test_invalid_integer_coercion(self): + self.assertEqual( + _matches_filter(self.sample_kv, ["max_num_tokens=four"]), + FilterResult.INVALID_FILTER) + + def test_invalid_boolean_coercion(self): + self.assertEqual(_matches_filter(self.sample_kv, ["use_bias=maybe"]), + FilterResult.INVALID_FILTER) + + def test_invalid_float_coercion(self): + self.assertEqual(_matches_filter(self.sample_kv, ["threshold=half"]), + FilterResult.INVALID_FILTER) + + def test_invalid_list_evaluation(self): + self.assertEqual( + _matches_filter(self.sample_kv, ["block_sizes=not_a_list"]), + FilterResult.INVALID_FILTER) + + +if __name__ == '__main__': + unittest.main() From a1d4117a9351566aabf38af00b4bb38fd14ff012 Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Sat, 6 Jun 2026 09:02:18 +0000 Subject: [PATCH 22/37] [skip ci] Update nightly support matrices (v6e/v7x) Signed-off-by: Buildkite Bot --- support_matrices/nightly/v6e/default/feature_support_matrix.csv | 2 +- support_matrices/nightly/v7x/default/feature_support_matrix.csv | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/support_matrices/nightly/v6e/default/feature_support_matrix.csv b/support_matrices/nightly/v6e/default/feature_support_matrix.csv index d085f597b5..f080152899 100644 --- a/support_matrices/nightly/v6e/default/feature_support_matrix.csv +++ b/support_matrices/nightly/v6e/default/feature_support_matrix.csv @@ -1,7 +1,7 @@ Feature,CorrectnessTest,PerformanceTest "Chunked Prefill",✅ Passing,✅ Passing "DCN-based P/D disaggregation",✅ Passing,✅ Passing -"KV Cache Offload",❌ Failing,❓ Untested +"KV Cache Offload",✅ Passing,✅ Passing "LoRA_Torch",✅ Passing,✅ Passing "Multimodal Inputs",✅ Passing,✅ Passing "Out-of-tree model support",✅ Passing,✅ Passing diff --git a/support_matrices/nightly/v7x/default/feature_support_matrix.csv b/support_matrices/nightly/v7x/default/feature_support_matrix.csv index 93ecef0f09..a9124ad39b 100644 --- a/support_matrices/nightly/v7x/default/feature_support_matrix.csv +++ b/support_matrices/nightly/v7x/default/feature_support_matrix.csv @@ -12,7 +12,7 @@ Feature,CorrectnessTest,PerformanceTest "Step Pooling (Embedding)",✅ Passing,❓ Untested "async scheduler",✅ Passing,✅ Passing "hybrid kv cache",✅ Passing,❓ Untested -"multi-host",✅ Passing,❓ Untested +"multi-host",❌ Failing,❓ Untested "runai_model_streamer_loader",✅ Passing,❓ Untested "sampling_params",✅ Passing,❓ Untested "structured_decoding",✅ Passing,❓ Untested From eff7e29f33dff445e79ecd1d372c448d90a7a26f Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Sat, 6 Jun 2026 12:01:39 +0000 Subject: [PATCH 23/37] [skip ci] Update nightly support matrices for flax_nnx (v6e/v7x) Signed-off-by: Buildkite Bot --- .../nightly/v7x/flax_nnx/feature_support_matrix.csv | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/support_matrices/nightly/v7x/flax_nnx/feature_support_matrix.csv b/support_matrices/nightly/v7x/flax_nnx/feature_support_matrix.csv index 93ecef0f09..a9124ad39b 100644 --- a/support_matrices/nightly/v7x/flax_nnx/feature_support_matrix.csv +++ b/support_matrices/nightly/v7x/flax_nnx/feature_support_matrix.csv @@ -12,7 +12,7 @@ Feature,CorrectnessTest,PerformanceTest "Step Pooling (Embedding)",✅ Passing,❓ Untested "async scheduler",✅ Passing,✅ Passing "hybrid kv cache",✅ Passing,❓ Untested -"multi-host",✅ Passing,❓ Untested +"multi-host",❌ Failing,❓ Untested "runai_model_streamer_loader",✅ Passing,❓ Untested "sampling_params",✅ Passing,❓ Untested "structured_decoding",✅ Passing,❓ Untested From 087e220f33d0165d1a124eb75ed8c38da55a0dd8 Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Sun, 7 Jun 2026 08:13:50 +0000 Subject: [PATCH 24/37] [skip ci] Update nightly support matrices (v6e/v7x) Signed-off-by: Buildkite Bot --- support_matrices/nightly/v7x/default/feature_support_matrix.csv | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/support_matrices/nightly/v7x/default/feature_support_matrix.csv b/support_matrices/nightly/v7x/default/feature_support_matrix.csv index a9124ad39b..93ecef0f09 100644 --- a/support_matrices/nightly/v7x/default/feature_support_matrix.csv +++ b/support_matrices/nightly/v7x/default/feature_support_matrix.csv @@ -12,7 +12,7 @@ Feature,CorrectnessTest,PerformanceTest "Step Pooling (Embedding)",✅ Passing,❓ Untested "async scheduler",✅ Passing,✅ Passing "hybrid kv cache",✅ Passing,❓ Untested -"multi-host",❌ Failing,❓ Untested +"multi-host",✅ Passing,❓ Untested "runai_model_streamer_loader",✅ Passing,❓ Untested "sampling_params",✅ Passing,❓ Untested "structured_decoding",✅ Passing,❓ Untested From ac548a044032cefccaddf8ef508925dc2cb6ac30 Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Sun, 7 Jun 2026 10:47:51 +0000 Subject: [PATCH 25/37] [skip ci] Update nightly support matrices for vllm (v6e/v7x) Signed-off-by: Buildkite Bot --- .../nightly/v6e/vllm/parallelism_support_matrix.csv | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv b/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv index a6cc9507ed..7abe71b054 100644 --- a/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv +++ b/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv @@ -1,6 +1,6 @@ Feature,Single-Host CorrectnessTest,Single-Host PerformanceTest,Multi-Host CorrectnessTest,Multi-Host PerformanceTest "CP",❓ Untested,❓ Untested,❓ Untested,❓ Untested -"DP",✅ Passing,✅ Passing,❓ Untested,❓ Untested +"DP",❌ Failing,❓ Untested,❓ Untested,❓ Untested "EP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "PP",✅ Passing,✅ Passing,✅ Passing,✅ Passing "SP",✅ Passing,❓ Untested,❓ Untested,❓ Untested From 0f665318f5e946f4933183225599c5ffb7c06866 Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Sun, 7 Jun 2026 10:49:34 +0000 Subject: [PATCH 26/37] [skip ci] Update nightly support matrices for flax_nnx (v6e/v7x) Signed-off-by: Buildkite Bot --- .../nightly/v7x/flax_nnx/feature_support_matrix.csv | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/support_matrices/nightly/v7x/flax_nnx/feature_support_matrix.csv b/support_matrices/nightly/v7x/flax_nnx/feature_support_matrix.csv index a9124ad39b..93ecef0f09 100644 --- a/support_matrices/nightly/v7x/flax_nnx/feature_support_matrix.csv +++ b/support_matrices/nightly/v7x/flax_nnx/feature_support_matrix.csv @@ -12,7 +12,7 @@ Feature,CorrectnessTest,PerformanceTest "Step Pooling (Embedding)",✅ Passing,❓ Untested "async scheduler",✅ Passing,✅ Passing "hybrid kv cache",✅ Passing,❓ Untested -"multi-host",❌ Failing,❓ Untested +"multi-host",✅ Passing,❓ Untested "runai_model_streamer_loader",✅ Passing,❓ Untested "sampling_params",✅ Passing,❓ Untested "structured_decoding",✅ Passing,❓ Untested From 3b2b06569130ddbf5c4aba8a08c27ee3e5397539 Mon Sep 17 00:00:00 2001 From: Kunjan Date: Sun, 7 Jun 2026 07:03:46 -0700 Subject: [PATCH 27/37] refactor into classes, vectorize ops (#2820) Signed-off-by: Kunjan Patel --- .../kernels/gdn/v2/compute_schedule_v2.py | 2 +- .../kernels/gdn/v2/recurrent_scan_impl.py | 930 ++++++++++++++++++ .../kernels/gdn/v2/recurrent_scan_v2.py | 888 +++-------------- 3 files changed, 1045 insertions(+), 775 deletions(-) create mode 100644 tpu_inference/kernels/gdn/v2/recurrent_scan_impl.py diff --git a/tpu_inference/kernels/gdn/v2/compute_schedule_v2.py b/tpu_inference/kernels/gdn/v2/compute_schedule_v2.py index 4be75ec7ac..c4d0a9d3cf 100644 --- a/tpu_inference/kernels/gdn/v2/compute_schedule_v2.py +++ b/tpu_inference/kernels/gdn/v2/compute_schedule_v2.py @@ -69,7 +69,7 @@ def compute_schedule_table_v2( # 1. Get each prefill sequence's effective start for chunkwise math # ========================================================================= r_idx = jnp.arange(num_seqs) - is_last_seq = r_idx == num_seqs - 1 + is_last_seq = r_idx == num_valid_seqs - 1 seq_start = query_start_loc[:-1] seq_end = query_start_loc[1:] num_tokens = query_start_loc[num_valid_seqs] diff --git a/tpu_inference/kernels/gdn/v2/recurrent_scan_impl.py b/tpu_inference/kernels/gdn/v2/recurrent_scan_impl.py new file mode 100644 index 0000000000..fe70e39095 --- /dev/null +++ b/tpu_inference/kernels/gdn/v2/recurrent_scan_impl.py @@ -0,0 +1,930 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Helper classes and processor implementation for Recurrent Scan Pallas kernel.""" + +# pylint: disable=invalid-name + +import dataclasses +from typing import Any + +import jax +import jax.numpy as jnp +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu + + +def l2_normalize(x, eps=1e-6): + rnorm = jax.lax.rsqrt(jnp.sum(x * x, axis=-1, keepdims=True) + eps) + return x * rnorm + + +# 1. Dataclasses for holding references to inputs/outputs and shared data. +# These are passed as arguments to the processor classes. + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class BranchRefs: + """Inputs/Outputs specific to a single execution branch (prefill or decode).""" + + qkv: Any + a_raw: Any + b_raw: Any + output: Any + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class SharedRefs: + """Inputs/Outputs refs shared between both branches(prefill and decode).""" + + a_log: Any + dt_bias: Any + recurrent_state_in: Any + recurrent_state_out: Any + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class PrefillScratchRefs: + """Scratch VMEM and semaphores allocated for prefill.""" + + scratch: Any + semaphore: Any + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class DecodeScratchRefs: + """Scratch VMEM and semaphores allocated for decode.""" + + state: Any + load: Any + store: Any + output: Any + read_semaphores: Any + write_semaphore: Any + + +# 2. Configuration Dataclasses (static) + + +@dataclasses.dataclass(frozen=True) +class ModelDims: + """Dimensions of the Recurrent Scan model configuration.""" + + n_kq: int + n_v: int + d_k: int + d_v: int + + @property + def key_dim(self) -> int: + return self.n_kq * self.d_k + + @property + def repeat_factor(self) -> int: + return self.n_v // self.n_kq + + +@dataclasses.dataclass(frozen=True) +class TilingConfig: + """Tiling dimensions for memory copy blocks.""" + + C: int + BT: int + sublanesize: int + + +@dataclasses.dataclass(frozen=True) +class ScanConfig: + """Configuration holding model dimensions and tiling options.""" + + model: ModelDims + tiling: TilingConfig + use_qk_norm_in_gdn: bool + decode_tokens: int + + +# 3. DMA Helper + + +class DMAHelper: + """Manages asynchronous state copies and double-buffering semaphores.""" + + def __init__(self, state_in, state_out, commit_scratch, semaphore): + self.state_in = state_in + self.state_out = state_out + self.commit_scratch = commit_scratch + self.sem = semaphore + # double or n buffering + self.has_multiple_slots = commit_scratch.shape[0] > 1 + + def build_copy_in(self, slot: int, state_idx: int): + target_slot = slot if self.has_multiple_slots else 0 + return pltpu.make_async_copy( + src_ref=self.state_in.at[pl.ds(state_idx, 1)], + dst_ref=self.commit_scratch.at[pl.ds(target_slot, 1)], + sem=self.sem.at[slot], + ) + + def commit_in(self, copy_op, slot: int, dst_ref, dst_slot: int): + target_slot = slot if self.has_multiple_slots else 0 + copy_op.wait() + dst_ref[dst_slot] = self.commit_scratch[target_slot].astype( + dst_ref.dtype) + + def copy_out(self, slot: int, state_idx: int, src_scratch): + target_slot = slot if self.has_multiple_slots else 0 + self.commit_scratch[target_slot] = src_scratch.astype( + self.commit_scratch.dtype) + copy_op = pltpu.make_async_copy( + src_ref=self.commit_scratch.at[pl.ds(target_slot, 1)], + dst_ref=self.state_out.at[pl.ds(state_idx, 1)], + sem=self.sem.at[slot], + ) + copy_op.start() + return copy_op + + def wait_out(self, slot: int, state_idx: int): + target_slot = slot if self.has_multiple_slots else 0 + copy_op = pltpu.make_async_copy( + src_ref=self.commit_scratch.at[pl.ds(target_slot, 1)], + dst_ref=self.state_out.at[pl.ds(state_idx, 1)], + sem=self.sem.at[slot], + ) + copy_op.wait() + + +# 4. Schedule step helper +COL_PREFILL_VALID = 0 +COL_PREFILL_OFFSET = 1 +COL_PREFILL_REQ_ID = 2 +COL_PREFILL_COUNT = 3 +COL_DECODE_VALID = 4 +COL_DECODE_OFFSET = 5 +COL_DECODE_REQ_ID = 6 +COL_DECODE_COUNT = 7 +COL_IS_LAST_CHUNK = 8 +COL_IS_FIRST_CHUNK = 9 +COL_IS_TRANSITION = 10 +COL_SUBLANE_REQ_IDS = 11 + + +class ScheduleStep: + """Unpacks and holds the scheduling metadata for the current step.""" + + def __init__(self, schedule_table, step): + self.step = step + self.schedule_table = schedule_table + self.prefill_valid = schedule_table[step, COL_PREFILL_VALID][...] + self.prefill_offset = schedule_table[step, COL_PREFILL_OFFSET][...] + self.prefill_req_id = schedule_table[step, COL_PREFILL_REQ_ID][...] + self.prefill_count = schedule_table[step, COL_PREFILL_COUNT][...] + + self.decode_valid = schedule_table[step, COL_DECODE_VALID][...] + self.decode_offset = schedule_table[step, COL_DECODE_OFFSET][...] + self.decode_req_id = schedule_table[step, COL_DECODE_REQ_ID][...] + self.decode_count = schedule_table[step, COL_DECODE_COUNT][...] + + self.is_last_chunk = schedule_table[step, COL_IS_LAST_CHUNK][...] + self.is_first_chunk = schedule_table[step, COL_IS_FIRST_CHUNK][...] + self.is_transition = schedule_table[step, COL_IS_TRANSITION][...] + + +# 4. Base Processor Class + + +class ScanProcessor: + """Base class for executing step calculations.""" + + def __init__( + self, + config: ScanConfig, + schedule: ScheduleStep, + state_indices, + has_initial_state, + ): + self.cfg = config + self.schedule = schedule + self.state_indices = state_indices + self.has_initial_state = has_initial_state + + +def invert_triangular_matrix(A, block_size=16): + """Inverts a unit lower triangular matrix A block-wise. + + Args: + A: Unit lower triangular matrix of shape (B, N, N). + block_size: Size of the blocks for Gaussian elimination. + + Returns: + Inverse of A, of shape (B, N, N). + """ + B, N, _ = A.shape + num_blocks = N // block_size + + def local_forward_sub(A_mat, b_mat): + x_list = [] + for i in range(block_size): + b_i = b_mat[:, i, :] + if i == 0: + x_i = b_i + else: + stacked_x = jnp.stack(x_list, axis=1) + all_prev_A = A_mat[:, i, :i] + prev_sum = jnp.sum(all_prev_A[..., None] * stacked_x, axis=1) + x_i = b_i - prev_sum + x_list.append(x_i) + return jnp.stack(x_list, axis=1) + + x_blocks = [] + for i in range(num_blocks): + start, end = i * block_size, (i + 1) * block_size + e_block = jnp.eye(N, dtype=A.dtype)[start:end, :] + e_block = jnp.broadcast_to(e_block, (B, block_size, N)) + + if i == 0: + target_b = e_block + else: + interaction_A = A[:, start:end, :start] + solved_x = jnp.concatenate(x_blocks, axis=1) + prev_sum = jnp.matmul(interaction_A, + solved_x, + precision=jax.lax.Precision.HIGHEST) + target_b = e_block - prev_sum + + local_A = A[:, start:end, start:end] + x_block = local_forward_sub(local_A, target_b) + x_blocks.append(x_block) + + return jnp.concatenate(x_blocks, axis=1) + + +class PrefillProcessor(ScanProcessor): + """Handles prefill step processing.""" + + def __init__( + self, + config: ScanConfig, + schedule: ScheduleStep, + state_indices, + has_initial_state, + refs: BranchRefs, + shared: SharedRefs, + scratch: PrefillScratchRefs, + dma: DMAHelper, + ): + super().__init__(config, schedule, state_indices, has_initial_state) + self.refs = refs + self.shared = shared + self.scratch = scratch + self.dma = dma + + def process(self): + is_trans = self.schedule.is_transition > 0 + jax.lax.cond( + is_trans, + lambda _: self._process_transition_prefill(), + lambda _: self._process_regular_prefill(), + operand=None, + ) + + def _process_regular_prefill(self): + """Processes a regular prefill step without transition boundary overlaps.""" + prefill_req_id = self.schedule.prefill_req_id + prefill_slot = prefill_req_id % 2 + init_has_init = self.has_initial_state[prefill_req_id][...] + init_state_idx = self.state_indices[prefill_req_id][...] + should_load_init = (self.schedule.is_first_chunk > 0) & (init_has_init + > 0) + should_zero_init = (self.schedule.is_first_chunk > 0) & (init_has_init + == 0) + + init_copy_op = self.dma.build_copy_in(prefill_slot, init_state_idx) + + @pl.when(should_load_init) + def _start_init_load(): + init_copy_op.start() + + @pl.when(should_zero_init) + def _zero_init_state(): + self.scratch.scratch[prefill_slot] = jnp.zeros( + (self.cfg.model.n_v, self.cfg.model.d_k, self.cfg.model.d_v), + dtype=self.scratch.scratch.dtype, + ) + + key_dim = self.cfg.model.key_dim + n_v = self.cfg.model.n_v + d_k = self.cfg.model.d_k + d_v = self.cfg.model.d_v + n_kq = self.cfg.model.n_kq + C = self.cfg.tiling.C + + qkv_chunk = self.refs.qkv[...] + qkv_chunk = jax.nn.silu(qkv_chunk) + q = qkv_chunk[:, :key_dim] + k = qkv_chunk[:, key_dim:2 * key_dim] + v = qkv_chunk[:, 2 * key_dim:] + + a_raw_chunk = self.refs.a_raw[...] + b_raw_chunk = self.refs.b_raw[...] + + a_raw_processed_T = a_raw_chunk[:, :n_v].T + b_raw_processed_T = b_raw_chunk[:, :n_v].T + + beta_T = jax.nn.sigmoid(b_raw_processed_T) + g_T = -jnp.exp( + # in jax 10.0.1 we can avoid the cast to float32, + # jax.errors.JaxRuntimeError: INTERNAL: Mosaic failed to compile TPU kernel: failed to legalize operation + # 'math.log1p': %7302 = "math.log1p"(%7295) <{fastmath = + # #arith.fastmath}> : (vector<8x128x2xbf16>) -> vector<8x128x2xbf16> + self.shared.a_log[...])[:, None] * jax.nn.softplus( + # same issue with the cast here + a_raw_processed_T + self.shared.dt_bias[...][:, None]) + g_T = jnp.maximum(g_T, -100.0) + + prefill_count = self.schedule.prefill_count + mask_float = (jnp.arange(C) < prefill_count).astype(q.dtype) + q = jnp.where(mask_float[:, None] > 0, q, 0.0) + k = jnp.where(mask_float[:, None] > 0, k, 0.0) + g_T = jnp.where(mask_float[None, :] > 0, g_T, 0.0) + v = jnp.where(mask_float[:, None] > 0, v, 0.0) + beta_T = jnp.where(mask_float[None, :] > 0, beta_T, 0.0) + + q = q.reshape(C, n_kq, d_k) + k = k.reshape(C, n_kq, d_k) + v = v.reshape(C, n_v, d_v) + + if self.cfg.use_qk_norm_in_gdn: + q = l2_normalize(q) + k = l2_normalize(k) + + # Note: fusing transpose with (vmatpush.xpose) made it slower, + # This has better instruction pipelining + q_T = q.transpose(1, 0, 2) # (n_kq, C, d_k) + k_T = k.transpose(1, 0, 2) # (n_kq, C, d_k) + v_T = v.transpose(1, 0, 2) # (n_v, C, d_v) + + repeat_factor = self.cfg.model.repeat_factor + if repeat_factor > 1: + q_T = jnp.repeat(q_T, repeat_factor, axis=0) + k_T = jnp.repeat(k_T, repeat_factor, axis=0) + + scale = d_k**-0.5 + q_T = q_T * scale + + g_cumsum_list = [] + current_sum = jnp.zeros((n_v, ), dtype=jnp.float32) + for i in range(C): + current_sum = current_sum + g_T[:, i] + g_cumsum_list.append(current_sum) + g_cumsum_T = jnp.stack(g_cumsum_list, axis=1) # shape (n_v, C) + exp_g = jnp.exp(g_cumsum_T)[..., None] # Precomputed for reuse + k_beta = k_T * beta_T[..., None] + + # Concatenate along sequence dimension: (n_v, 2 * C, d_k) + kbeta_q = jnp.concatenate([k_beta, q_T], axis=1) + # Batch is n_v (axis 0), contract is d_k (axis 2). + # Output shape: (n_v, 2 * C, C) + S_both = jax.lax.dot_general( + kbeta_q, + k_T, + (((2, ), (2, )), ((0, ), (0, ))), + preferred_element_type=jnp.float32, + ) + S = S_both[:, :C, :] + S_q = S_both[:, C:, :] + + g_diff = g_cumsum_T[..., :, None] - g_cumsum_T[..., None, :] + i_idx = jnp.arange(C)[:, None] + j_idx = jnp.arange(C)[None, :] + mask_float = (i_idx > j_idx).astype(jnp.float32) + + g_diff_safe = jnp.minimum(g_diff, 0.0) + S = jnp.where(mask_float[None, :, :] > 0, S * jnp.exp(g_diff_safe), + 0.0) + + mask_float_q = (i_idx >= j_idx).astype(jnp.float32) + g_diff_Sq = g_diff_safe * mask_float_q[None, ...] + ( + 1.0 - mask_float_q[None, ...]) * (-1e30) + S_q = S_q * jnp.exp(g_diff_Sq) + S_q = S_q * mask_float_q[None, ...] + + I_plus_S = jnp.eye(C, dtype=jnp.float32)[None, ...] + S + A_inv = invert_triangular_matrix(I_plus_S, block_size=16) + + v_beta = v_T * beta_T[..., None] + k_beta_g = k_beta * exp_g + vk_in = jnp.concatenate( + [ + v_beta, + k_beta_g, + ], + axis=2, + ) # (n_v, C, d_v + d_k) + uw = jax.lax.dot_general( + A_inv, + vk_in, + (((2, ), (1, )), ((0, ), (0, ))), + precision=jax.lax.Precision.HIGHEST, + ) # Output shape: (n_v, C, d_v + d_k) + u = uw[..., :d_v] + w = uw[..., d_v:] + + q_g = q_T * exp_g # (n_v, C, d_k) + + @pl.when(should_load_init) + def _finish_init_load(): + self.dma.commit_in(init_copy_op, prefill_slot, + self.scratch.scratch, prefill_slot) + + current_state = self.scratch.scratch[prefill_slot] # (n_v, d_k, d_v) + + qw = jnp.concatenate([q_g, w], axis=1) # (n_v, 2 * C, d_k) + comb = jax.lax.dot_general( + qw, + current_state.astype(jnp.float32), + (((2, ), (1, )), ((0, ), (0, ))), + precision=jax.lax.Precision.DEFAULT, + ) # Output shape: (n_v, 2 * C, d_v) + attn_inter, v_prime = jnp.split(comb, 2, axis=1) + + v_new = u - v_prime + term2 = jnp.matmul(S_q, v_new, precision=jax.lax.Precision.HIGHEST) + o_c = attn_inter + term2 # (n_v, C, d_v) + + g_i_last_exp = exp_g[:, -1, None] + g_diff_exp_state = jnp.exp(g_cumsum_T[..., -1, None] - + g_cumsum_T)[..., None] + k_i_g_diff = k_T * g_diff_exp_state + update_term = jax.lax.dot_general( + k_i_g_diff, + v_new, + (((1, ), (1, )), ((0, ), (0, ))), + precision=jax.lax.Precision.DEFAULT, + ) # Output shape: (n_v, d_k, d_v) + h_new = current_state * g_i_last_exp + update_term + + self.scratch.scratch[prefill_slot] = h_new.astype( + self.scratch.scratch.dtype) + + store_state_idx = self.state_indices[prefill_req_id][...] + + @pl.when(self.schedule.is_last_chunk > 0) + def store_state(): + copy_op = self.dma.copy_out(prefill_slot, store_state_idx, + self.scratch.scratch[prefill_slot]) + copy_op.wait() + + o_c_tr = o_c.transpose(1, 0, 2) + o_c_flat = o_c_tr.reshape(C, n_v * d_v) + + mask_float = (jnp.arange(C) < prefill_count).astype(o_c_flat.dtype) + o_c_flat_masked = o_c_flat * mask_float[:, None] + self.refs.output[...] = o_c_flat_masked.astype(self.refs.output.dtype) + + def _process_transition_prefill(self): + """Processes a transition prefill step with sublane stitching.""" + C_trans = self.cfg.tiling.sublanesize + key_dim = self.cfg.model.key_dim + n_v = self.cfg.model.n_v + d_k = self.cfg.model.d_k + d_v = self.cfg.model.d_v + n_kq = self.cfg.model.n_kq + + first_req_id = self.schedule.schedule_table[self.schedule.step, + COL_SUBLANE_REQ_IDS][...] + first_is_first = self.schedule.schedule_table[self.schedule.step, + COL_SUBLANE_REQ_IDS + + C_trans][...] + first_slot = first_req_id % 2 + first_has_init = self.has_initial_state[first_req_id][...] + should_load_first = (first_is_first > 0) & (first_has_init > 0) + first_state_idx = self.state_indices[first_req_id][...] + + first_copy_op = self.dma.build_copy_in(first_slot, first_state_idx) + + @pl.when(should_load_first) + def _start_first_load(): + first_copy_op.start() + + qkv_chunk = self.refs.qkv[:C_trans, :] + qkv_chunk = jax.nn.silu(qkv_chunk) + q = qkv_chunk[:, :key_dim] + k = qkv_chunk[:, key_dim:2 * key_dim] + v = qkv_chunk[:, 2 * key_dim:] + + a_raw_chunk = self.refs.a_raw[...] + b_raw_chunk = self.refs.b_raw[...] + + a_raw_processed_T = a_raw_chunk[:C_trans, :n_v].T + b_raw_processed_T = b_raw_chunk[:C_trans, :n_v].T + + beta_chunk_T = jax.nn.sigmoid(b_raw_processed_T) + g_chunk_T = -jnp.exp( + self.shared.a_log[...])[:, None] * jax.nn.softplus( + a_raw_processed_T + self.shared.dt_bias[...][:, None]) + g_chunk_T = jnp.maximum(g_chunk_T, -100.0) + + q = q.reshape(C_trans, n_kq, d_k) + k = k.reshape(C_trans, n_kq, d_k) + v = v.reshape(C_trans, n_v, d_v) + + if self.cfg.use_qk_norm_in_gdn: + q = l2_normalize(q) + k = l2_normalize(k) + + repeat_factor = self.cfg.model.repeat_factor + if repeat_factor > 1: + q = jnp.repeat(q, repeat_factor, axis=1) + k = jnp.repeat(k, repeat_factor, axis=1) + + scale = d_k**-0.5 + q = q * scale + + @pl.when((first_is_first > 0) & (first_has_init == 0)) + def _zero_first_slot(): + self.scratch.scratch[first_slot] = jnp.zeros( + (n_v, d_k, d_v), dtype=self.scratch.scratch.dtype) + + @pl.when(should_load_first) + def _finish_first_load(): + self.dma.commit_in(first_copy_op, first_slot, self.scratch.scratch, + first_slot) + + h = self.scratch.scratch[first_slot] + current_r = first_req_id + sequence_valid = True + exp_g_chunk_T = jnp.exp(g_chunk_T) + + for i in range(C_trans): + t_req = self.schedule.schedule_table[self.schedule.step, + 11 + i][...] + t_is_first = self.schedule.schedule_table[self.schedule.step, + 11 + C_trans + i][...] + t_is_last = self.schedule.schedule_table[self.schedule.step, + 11 + 2 * C_trans + i][...] + + is_new_seq = t_req != current_r + sequence_valid = jnp.where(is_new_seq, True, sequence_valid) + + is_decode_token = t_req < self.cfg.decode_tokens + sequence_valid = jnp.where(is_decode_token, False, sequence_valid) + + c_slot = current_r % 2 + self.scratch.scratch[c_slot] = h + + def do_write(c_slot=c_slot, current_r=current_r, h=h): + state_idx = self.state_indices[current_r][...] + copy_op = self.dma.copy_out(c_slot, state_idx, h) + copy_op.wait() + return None + + is_current_r_prefill = current_r >= self.cfg.decode_tokens + should_write = is_current_r_prefill & is_new_seq + jax.lax.cond(should_write, do_write, lambda: None) + + t_slot = t_req % 2 + t_has_init = self.has_initial_state[t_req][...] + + def load_t_state(t_slot=t_slot, t_req=t_req): + state_idx = self.state_indices[t_req][...] + copy_op = self.dma.build_copy_in(t_slot, state_idx) + copy_op.start() + self.dma.commit_in(copy_op, t_slot, self.scratch.scratch, + t_slot) + + should_load_t = (t_is_first > 0) & (t_has_init > 0) + jax.lax.cond(should_load_t, load_t_state, lambda: None) + + should_zero = (t_is_first > 0) & (t_has_init == 0) + + def zero_t_slot(t_slot=t_slot, n_v=n_v, d_k=d_k, d_v=d_v): + self.scratch.scratch[t_slot] = jnp.zeros( + (n_v, d_k, d_v), dtype=self.scratch.scratch.dtype) + + jax.lax.cond(should_zero, zero_t_slot, lambda: None) + + h = self.scratch.scratch[t_slot] + current_r = t_req + + k_i = k[i, :, :] + v_i = v[i, :, :] + beta_i = beta_chunk_T[:, i] + q_i = q[i, :, :] + + decay = exp_g_chunk_T[:, i][..., None] + + k_state = jnp.sum(k_i[..., None] * h, axis=1) + v_diff = v_i - decay * k_state + v_new = beta_i[:, None] * v_diff + + q_state = jnp.sum(q_i[..., None] * h, axis=1) + q_k = jnp.sum(q_i * k_i, axis=-1, keepdims=True) + + out_i = decay * q_state + q_k * v_new + + k_v_new = k_i[..., None] * v_new[:, None, :] + h_new = h * decay[..., None] + k_v_new + + h = jnp.where(sequence_valid, h_new, h) + out_i = jnp.where(sequence_valid, out_i, 0.0) + + sequence_valid = jnp.where(t_is_last > 0, False, sequence_valid) + + self.refs.output[i, :] = out_i.reshape(n_v * d_v).astype( + self.refs.output.dtype) + + final_slot = current_r % 2 + self.scratch.scratch[final_slot] = h + + is_current_r_prefill = current_r >= self.cfg.decode_tokens + + @pl.when(is_current_r_prefill) + def do_final_write(): + state_idx = self.state_indices[current_r][...] + copy_op = self.dma.copy_out(final_slot, state_idx, h) + copy_op.wait() + return None + + +class DecodeProcessor(ScanProcessor): + """Handles batch decode step processing using double-buffering logic.""" + + def __init__( + self, + config: ScanConfig, + schedule: ScheduleStep, + state_indices, + has_initial_state, + refs: BranchRefs, + shared: SharedRefs, + scratch: DecodeScratchRefs, + dma: DMAHelper, + ): + super().__init__(config, schedule, state_indices, has_initial_state) + self.refs = refs + self.shared = shared + self.scratch = scratch + self.dma = dma + + def get_target_idx(self, b): + safe_req_id = jnp.minimum(self.schedule.decode_req_id + b, + self.state_indices.shape[0] - 1) + return self.state_indices[safe_req_id][...] + + def process(self): + """Processes decode steps in blocks.""" + decode_count = self.schedule.decode_count + BT = self.cfg.tiling.BT + n_v = self.cfg.model.n_v + d_k = self.cfg.model.d_k + d_v = self.cfg.model.d_v + n_kq = self.cfg.model.n_kq + key_dim = self.cfg.model.key_dim + repeat_factor = self.cfg.model.repeat_factor + use_qk_norm_in_gdn = self.cfg.use_qk_norm_in_gdn + exp_a_log = jnp.exp(self.shared.a_log[...].astype(jnp.float32)) + dt_bias_f32 = self.shared.dt_bias[...].astype(jnp.float32) + + # Pre-loop: kick off async loads for iters 0 and 1. + @pl.when(decode_count >= 1) + def _preload_slot_0(): + tgt = self.get_target_idx(0) + op = pltpu.make_async_copy( + src_ref=self.shared.recurrent_state_in.at[pl.ds(tgt, 1)], + dst_ref=self.scratch.load.at[pl.ds(0, 1)], + sem=self.scratch.read_semaphores.at[0], + ) + op.start() + + @pl.when(decode_count >= 2) + def _preload_slot_1(): + tgt = self.get_target_idx(1) + op = pltpu.make_async_copy( + src_ref=self.shared.recurrent_state_in.at[pl.ds(tgt, 1)], + dst_ref=self.scratch.load.at[pl.ds(1, 1)], + sem=self.scratch.read_semaphores.at[1], + ) + op.start() + + def process_decode_step(b, store_inflight): + s0_inflight, s1_inflight = store_inflight + is_valid = b < decode_count + slot = b % 2 + using_slot_0 = slot == 0 + cur_slot_inflight = jax.lax.select(using_slot_0, s0_inflight, + s1_inflight) + + @pl.when(is_valid) + def do_work(): + # Wait for THIS iter's load. + wait_load = pltpu.make_async_copy( + src_ref=self.shared.recurrent_state_in.at[pl.ds(0, 1)], + dst_ref=self.scratch.load.at[pl.ds(slot, 1)], + sem=self.scratch.read_semaphores.at[slot], + ) + wait_load.wait() + + self.scratch.state[pl.ds(0, 1)] = self.scratch.load[pl.ds( + slot, 1)][...] + + # Prefetch load for iter b+2. + next_b = b + 2 + + @pl.when(next_b < decode_count) + def _prefetch_next_load(): + next_tgt = self.get_target_idx(next_b) + op = pltpu.make_async_copy( + src_ref=self.shared.recurrent_state_in.at[pl.ds( + next_tgt, 1)], + dst_ref=self.scratch.load.at[pl.ds(slot, 1)], + sem=self.scratch.read_semaphores.at[slot], + ) + op.start() + + target_idx = self.get_target_idx(b) + + sublanesize = self.cfg.tiling.sublanesize + b_aligned = (b // sublanesize) * sublanesize + + qkv_block_data = self.refs.qkv[ + pl.ds(b_aligned, sublanesize), :].astype(jnp.float32) + mask = (jnp.arange(sublanesize) == (b % sublanesize)).astype( + qkv_block_data.dtype)[:, None] + qkv_row = jnp.sum(qkv_block_data * mask, axis=0, keepdims=True) + + # Fused SiLU + qkv_row = jax.nn.silu(qkv_row) + q = qkv_row[:, :key_dim].reshape(n_kq, d_k) + k = qkv_row[:, key_dim:2 * key_dim].reshape(n_kq, d_k) + v = qkv_row[:, 2 * key_dim:].reshape(n_v, d_v) + + if use_qk_norm_in_gdn: + q = l2_normalize(q) + k = l2_normalize(k) + + # Head repetition + if repeat_factor > 1: + q = jnp.repeat(q, repeat_factor, axis=0) + k = jnp.repeat(k, repeat_factor, axis=0) + + scale = d_k**-0.5 + q = q * scale + + g_block_new = self.refs.a_raw[pl.ds(b_aligned, sublanesize), :] + beta_block_new = self.refs.b_raw[ + pl.ds(b_aligned, sublanesize), :] + + mask_new = (jnp.arange(sublanesize) == ( + b % sublanesize)).astype(g_block_new.dtype)[:, None] + + a_raw_new = jnp.sum(g_block_new * mask_new, + axis=0, + keepdims=True)[0, :n_v].astype(jnp.float32) + b_raw_new = jnp.sum(beta_block_new * mask_new, + axis=0, + keepdims=True)[0, :n_v].astype(jnp.float32) + + # Compute gate + curr_beta = jax.nn.sigmoid(b_raw_new) + curr_g = -exp_a_log * jax.nn.softplus(a_raw_new + dt_bias_f32) + curr_g = jnp.maximum(curr_g, -100.0) + decay = jnp.exp(curr_g) + + current_state = self.scratch.state[0] + + # 1. Batched dot product: k @ state -> (n_v, d_v) + k_state = jax.lax.dot_general( + k.reshape(n_v, 1, d_k), + current_state, + (((2, ), (1, )), ((0, ), (0, ))), + preferred_element_type=jnp.float32, + ).reshape(n_v, d_v) + + decay_k_state = jnp.where( + jnp.isinf(k_state), + 0.0, + decay[:, None] * k_state, + ) + v_diff = v - decay_k_state + v_new = curr_beta[:, None] * v_diff + + # 2. Batched dot product: q @ state -> (n_v, d_v) + q_state = jax.lax.dot_general( + q.reshape(n_v, 1, d_k), + current_state, + (((2, ), (1, )), ((0, ), (0, ))), + preferred_element_type=jnp.float32, + ).reshape(n_v, d_v) + + q_k = jnp.sum( + q * k, + axis=-1, + keepdims=True, + ) + + decay_q_state = jnp.where( + jnp.isinf(q_state), + 0.0, + decay[:, None] * q_state, + ) + out_step = decay_q_state + q_k * v_new + + # 3. Outer product and decay update for state -> (n_v, d_k, d_v) + decay_state = jnp.where( + jnp.isinf(current_state), + 0.0, + current_state * decay[:, None, None], + ) + k_v_new = k[:, :, None] * v_new[:, None, :] + new_state = decay_state + k_v_new + self.scratch.store[slot] = new_state.astype( + self.scratch.store.dtype) + + # Accumulate output in scratchpad + current_output = self.scratch.output[...] + mask = (jnp.arange(BT) == b).astype(current_output.dtype)[:, + None] + new_output = jnp.where( + mask, + out_step.reshape(1, + n_v * d_v).astype(current_output.dtype), + current_output, + ) + self.scratch.output[...] = new_output + + # Async store. Before writing to decode_store_scratch[slot], + # wait for the previous same-slot store DMA (from iter b-2) + @pl.when(cur_slot_inflight > 0) + def _wait_same_slot_store(): + copy_op = pltpu.make_async_copy( + src_ref=self.scratch.store.at[pl.ds(slot, 1)], + dst_ref=self.shared.recurrent_state_out.at[pl.ds(0, + 1)], + sem=self.scratch.write_semaphore.at[slot], + ) + copy_op.wait() + + copy_op = pltpu.make_async_copy( + src_ref=self.scratch.store.at[pl.ds(slot, 1)], + dst_ref=self.shared.recurrent_state_out.at[pl.ds( + target_idx, 1)], + sem=self.scratch.write_semaphore.at[slot], + ) + copy_op.start() + + next_s0_inflight = jax.lax.select( + is_valid & using_slot_0, + jnp.int32(1), + s0_inflight, + ) + next_s1_inflight = jax.lax.select( + is_valid & (~using_slot_0), + jnp.int32(1), + s1_inflight, + ) + return (next_s0_inflight, next_s1_inflight) + + final_s0_inflight, final_s1_inflight = jax.lax.fori_loop( + 0, + BT, + process_decode_step, + (jnp.int32(0), jnp.int32(0)), + ) + + # Drain any remaining async store DMAs + @pl.when(final_s0_inflight > 0) + def _drain_slot_0(): + temp_desc = pltpu.make_async_copy( + src_ref=self.scratch.store.at[pl.ds(0, 1)], + dst_ref=self.shared.recurrent_state_out.at[pl.ds(0, 1)], + sem=self.scratch.write_semaphore.at[0], + ) + temp_desc.wait() + + @pl.when(final_s1_inflight > 0) + def _drain_slot_1(): + temp_desc = pltpu.make_async_copy( + src_ref=self.scratch.store.at[pl.ds(1, 1)], + dst_ref=self.shared.recurrent_state_out.at[pl.ds(0, 1)], + sem=self.scratch.write_semaphore.at[1], + ) + temp_desc.wait() + + mask = (jnp.arange(BT) + < decode_count).astype(self.scratch.output.dtype)[:, None] + decode_output_scratch_masked = self.scratch.output[...] * mask + self.refs.output[...] = decode_output_scratch_masked diff --git a/tpu_inference/kernels/gdn/v2/recurrent_scan_v2.py b/tpu_inference/kernels/gdn/v2/recurrent_scan_v2.py index a4f84e98a7..3756a24219 100644 --- a/tpu_inference/kernels/gdn/v2/recurrent_scan_v2.py +++ b/tpu_inference/kernels/gdn/v2/recurrent_scan_v2.py @@ -11,6 +11,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""Pallas kernel for GDN recurrent scan.""" + +# pylint: disable=invalid-name import functools @@ -21,34 +24,7 @@ from tpu_inference.kernels.gdn.v2 import \ compute_schedule_v2 as compute_schedule_table_v2 - -# def invert_triangular_matrix(A, block_size=None): -# """Inverts a unit lower triangular matrix A using Neumann doubling. - -# Algorithm: Neumann doubling. For L strictly lower triangular of size N x N, -# since L^N = 0 we have: -# (I + L)^{-1} = (I - L)(I + L^2)(I + L^4) ... (I + L^(N/2)) - -# Args: -# A: Unit lower triangular matrix of shape (B, N, N). -# block_size: Size of the blocks for Gaussian elimination (unused). - -# Returns: -# Inverse of A, of shape (B, N, N). - -# For N=128 this is exactly 6 iterations = 12 matmuls, all (B, 128, 128). -# """ -# B, N, _ = A.shape -# num_iters = max(1, (N - 1).bit_length() - 1) -# in_dtype = A.dtype -# A_f32 = A.astype(jnp.float32) -# L = jnp.tril(A_f32, k=-1) -# eye = jnp.broadcast_to(jnp.eye(N, dtype=jnp.float32), (B, N, N)) -# Y = eye - L -# for _ in range(num_iters): -# L = jnp.matmul(L, L, precision=jax.lax.Precision.HIGHEST) -# Y = Y + jnp.matmul(Y, L, precision=jax.lax.Precision.HIGHEST) -# return Y.astype(in_dtype) +from tpu_inference.kernels.gdn.v2 import recurrent_scan_impl def invert_triangular_matrix(A, block_size=16): @@ -184,759 +160,126 @@ def inner_kernel( """ step = pl.program_id(0) - # READ table + # Instantiate helper config structures + model_dims = recurrent_scan_impl.ModelDims(n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v) + tiling_cfg = recurrent_scan_impl.TilingConfig(C=C, + BT=BT, + sublanesize=sublanesize) + config = recurrent_scan_impl.ScanConfig( + model=model_dims, + tiling=tiling_cfg, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + decode_tokens=decode_tokens, + ) - prefill_valid = schedule_table[step, 0][...] - prefill_req_id = schedule_table[step, 2][...] + schedule = recurrent_scan_impl.ScheduleStep(schedule_table, step) - decode_valid = schedule_table[step, 4][...] - decode_offset = schedule_table[step, 5][...] - decode_req_id = schedule_table[step, 6][...] - decode_count = schedule_table[step, 7][...] + shared_refs = recurrent_scan_impl.SharedRefs( + a_log=a_log_ref, + dt_bias=dt_bias_ref, + recurrent_state_in=recurrent_state_in, + recurrent_state_out=recurrent_state_out, + ) - prefill_offset = schedule_table[step, 1][...] - is_transition = schedule_table[step, 10][...] + prefill_refs = recurrent_scan_impl.BranchRefs( + qkv=prefill_qkv_ref, + a_raw=prefill_a_raw_ref, + b_raw=prefill_b_raw_ref, + output=prefill_output_ref, + ) + prefill_scratch_refs = recurrent_scan_impl.PrefillScratchRefs( + scratch=prefill_scratch, + semaphore=prefill_semaphore, + ) + prefill_dma = recurrent_scan_impl.DMAHelper( + state_in=recurrent_state_in, + state_out=recurrent_state_out, + commit_scratch=state_commit_scratch, + semaphore=prefill_semaphore, + ) - is_last_chunk = schedule_table[step, 8][...] - is_first_chunk = schedule_table[step, 9][...] + prefill_processor = recurrent_scan_impl.PrefillProcessor( + config=config, + schedule=schedule, + state_indices=state_indices, + has_initial_state=has_initial_state, + refs=prefill_refs, + shared=shared_refs, + scratch=prefill_scratch_refs, + dma=prefill_dma, + ) - def l2_normalize(x, eps=1e-6): - norm = jnp.sqrt(jnp.sum(x * x, axis=-1, keepdims=True) + eps) - return x / norm + decode_refs = recurrent_scan_impl.BranchRefs( + qkv=decode_qkv_ref, + a_raw=decode_a_raw_ref, + b_raw=decode_b_raw_ref, + output=decode_output_ref, + ) + decode_scratch = recurrent_scan_impl.DecodeScratchRefs( + state=decode_state_scratch, + load=decode_load_scratch, + store=decode_store_scratch, + output=decode_output_scratch, + read_semaphores=decode_read_semaphores, + write_semaphore=decode_write_semaphore, + ) + decode_dma_in = recurrent_scan_impl.DMAHelper( + state_in=recurrent_state_in, + state_out=recurrent_state_out, + commit_scratch=decode_load_scratch, + semaphore=decode_read_semaphores, + ) + decode_processor = recurrent_scan_impl.DecodeProcessor( + config=config, + schedule=schedule, + state_indices=state_indices, + has_initial_state=has_initial_state, + refs=decode_refs, + shared=shared_refs, + scratch=decode_scratch, + dma=decode_dma_in, + ) + + # READ table + + prefill_valid = schedule_table[step, + recurrent_scan_impl.COL_PREFILL_VALID][...] + decode_valid = schedule_table[step, + recurrent_scan_impl.COL_DECODE_VALID][...] + decode_offset = schedule_table[step, + recurrent_scan_impl.COL_DECODE_OFFSET][...] + prefill_offset = schedule_table[ + step, recurrent_scan_impl.COL_PREFILL_OFFSET][...] + is_transition = schedule_table[step, + recurrent_scan_impl.COL_IS_TRANSITION][...] # 2. Decode Branch @pl.when(decode_valid > 0) def decode_wrapper(): - - def get_target_idx(b): - safe_req_id = jnp.minimum(decode_req_id + b, - state_indices.shape[0] - 1) - return state_indices[safe_req_id][...] - - # Pre-loop: kick off async loads for iters 0 and 1. - # iter b consumes the load that lands in decode_load_scratch[b % 2]. - # Subsequent loads (iter b+2 for each iter b) are issued from inside - # the loop as prefetches. - @pl.when(decode_count >= 1) - def _preload_slot_0(): - tgt = get_target_idx(0) - op = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(tgt, 1)], - dst_ref=decode_load_scratch.at[pl.ds(0, 1)], - sem=decode_read_semaphores.at[0], - ) - op.start() - - @pl.when(decode_count >= 2) - def _preload_slot_1(): - tgt = get_target_idx(1) - op = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(tgt, 1)], - dst_ref=decode_load_scratch.at[pl.ds(1, 1)], - sem=decode_read_semaphores.at[1], - ) - op.start() - - def process_decode(b, store_inflight): - # store_inflight: tuple (s0_inflight, s1_inflight) of int32 scalars. - # s{n}_inflight == 1 iff slot n has an in-flight async store DMA. - s0_inflight, s1_inflight = store_inflight - is_valid = b < decode_count - slot = b % 2 - using_slot_0 = slot == 0 - cur_slot_inflight = jax.lax.select(using_slot_0, s0_inflight, - s1_inflight) - - @pl.when(is_valid) - def do_work(): - # Wait for THIS iter's load (issued by preload or by iter b-2 prefetch). - wait_load = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(0, 1)], - dst_ref=decode_load_scratch.at[pl.ds(slot, 1)], - sem=decode_read_semaphores.at[slot], - ) - wait_load.wait() - - # Safe-copy of loaded state. Isolates compute from the - # prefetch DMA below that writes the same slot of - # decode_load_scratch concurrently. bf16 -> bf16, no cast. - decode_state_scratch[pl.ds(0, 1)] = decode_load_scratch[pl.ds( - slot, 1)][...] - - # Prefetch load for iter b+2 (same slot, since (b+2) % 2 == b % 2). - # This DMA overlaps with the compute below. - next_b = b + 2 - - @pl.when(next_b < decode_count) - def _prefetch_next_load(): - next_tgt = get_target_idx(next_b) - op = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(next_tgt, 1)], - dst_ref=decode_load_scratch.at[pl.ds(slot, 1)], - sem=decode_read_semaphores.at[slot], - ) - op.start() - - target_idx = get_target_idx(b) - - key_dim = n_kq * d_k - b_aligned = (b // sublanesize) * sublanesize - # Workaround: Upcast to fp32 to avoid NaNs - qkv_block_data = decode_qkv_ref[ - pl.ds(b_aligned, sublanesize), :].astype(jnp.float32) - mask = (jnp.arange(sublanesize) == (b % sublanesize)).astype( - qkv_block_data.dtype)[:, None] - qkv_row = jnp.sum(qkv_block_data * mask, axis=0, keepdims=True) - # Fused SiLU - qkv_row = jax.nn.silu(qkv_row) - q = qkv_row[:, :key_dim].reshape(n_kq, d_k) - k = qkv_row[:, key_dim:2 * key_dim].reshape(n_kq, d_k) - v = qkv_row[:, 2 * key_dim:].reshape(n_v, d_v) - - if use_qk_norm_in_gdn: - q = l2_normalize(q) - k = l2_normalize(k) - - # Head repetition - repeat_factor = n_v // n_kq - if repeat_factor > 1: - q = jnp.repeat(q, repeat_factor, axis=0) - k = jnp.repeat(k, repeat_factor, axis=0) - - scale = d_k**-0.5 - q = q * scale - - b_aligned = (b // sublanesize) * sublanesize - - g_block_new = decode_a_raw_ref[ - pl.ds(b_aligned, sublanesize), :] - beta_block_new = decode_b_raw_ref[ - pl.ds(b_aligned, sublanesize), :] - - mask_new = (jnp.arange(sublanesize) == ( - b % sublanesize)).astype(g_block_new.dtype)[:, None] - - curr_g_slice_new = jnp.sum(g_block_new * mask_new, - axis=0, - keepdims=True) - curr_beta_slice_new = jnp.sum(beta_block_new * mask_new, - axis=0, - keepdims=True) - - a_raw_new = curr_g_slice_new[:, :n_v].reshape(n_v).astype( - jnp.float32) - b_raw_new = (curr_beta_slice_new[:, :n_v].reshape(n_v).astype( - jnp.float32)) - - # Compute gate - curr_beta = jax.nn.sigmoid(b_raw_new) - curr_g = -jnp.exp(a_log_ref[...].astype( - jnp.float32)) * jax.nn.softplus( - a_raw_new + dt_bias_ref[...].astype(jnp.float32)) - curr_g = jnp.maximum(curr_g, -100.0) - decay = jnp.exp(curr_g) - - current_state = decode_state_scratch[0] - - # TODO: compare MXU vs VPU, MXU doesn't support FP32, VPU does - # (n_v, d_k, 1) * (n_v, 1, d_v) -> (n_v, d_k, d_v) - out_list = [] - new_state_list = [] - for h in range(n_v): - q_h = q[h:h + 1, :] # (1, d_k) - k_h = k[h:h + 1, :] # (1, d_k) - v_h = v[h:h + 1, :] # (1, d_v) - - state_h = current_state[h].astype( - jnp.float32) # (d_k, d_v) - - k_state_h = pl.dot( - k_h, state_h, - precision=jax.lax.Precision.HIGHEST) # (1, d_v) - - # v_diff_h = v_h - decay[h].astype(jnp.float32) * k_state_h - decay_k_state = jnp.where( - jnp.isinf(k_state_h), - 0.0, - decay[h].astype(jnp.float32) * k_state_h, - ) - v_diff_h = v_h - decay_k_state - v_new_h = curr_beta[h].astype(jnp.float32) * v_diff_h - - q_state_h = pl.dot( - q_h, state_h, - precision=jax.lax.Precision.HIGHEST) # (1, d_v) - - q_k_h = jnp.sum(q_h * k_h, axis=-1, - keepdims=True) # (1, 1) - - # Defensive code to handle NaNs and infs in state, - # Saw similar issue while trying newton schulz - # which can happen due to large decay or long sequences. - # TODO: analyze perf impact and risk of removing this. - decay_q_state = jnp.where(jnp.isinf(q_state_h), 0.0, - decay[h] * q_state_h) - out_h = decay_q_state + q_k_h * v_new_h - out_list.append(out_h) - - k_v_new_h = pl.dot(k_h, - v_new_h, - trans_a=True, - precision=jax.lax.Precision.HIGHEST - ) # (d_k, 1) @ (1, d_v) -> (d_k, d_v) - # Defensive code to handle NaNs and infs in state, - # which can happen due to large decay or long sequences. - # In such cases, we reset the state contribution to zero and rely solely on the new value - # TODO: analyze perf impact and risk of removing this. - decay_state = jnp.where(jnp.isinf(state_h), 0.0, - state_h * decay[h]) - new_state_h = decay_state + k_v_new_h - new_state_list.append(new_state_h) - - out = jnp.concatenate(out_list, axis=0) # (n_v, d_v) - new_state = jnp.stack(new_state_list, - axis=0) # (n_v, d_k, d_v) - - # Accumulate output in scratchpad - current_output = decode_output_scratch[...] - mask = (jnp.arange(BT) == b).astype(current_output.dtype)[:, - None] - new_output = jnp.where( - mask, - out.reshape(1, n_v * d_v), - current_output, - ) - decode_output_scratch[...] = new_output.astype( - current_output.dtype) - - # Async store. Before writing to decode_store_scratch[slot], - # wait for the previous same-slot store DMA (from iter b-2) - # so we don't clobber a buffer that's still being read. - @pl.when(cur_slot_inflight > 0) - def _wait_same_slot_store(): - temp_desc = pltpu.make_async_copy( - src_ref=decode_store_scratch.at[pl.ds(slot, 1)], - dst_ref=recurrent_state_out.at[pl.ds(0, 1)], - sem=decode_write_semaphore.at[slot], - ) - temp_desc.wait() - - decode_store_scratch[slot] = new_state.astype( - decode_store_scratch.dtype) - copy_op = pltpu.make_async_copy( - src_ref=decode_store_scratch.at[pl.ds(slot, 1)], - dst_ref=recurrent_state_out.at[pl.ds(target_idx, 1)], - sem=decode_write_semaphore.at[slot], - ) - copy_op.start() - # No wait — drained after fori_loop. - - # Update carry: this iter marked its slot as having an in-flight - # store iff is_valid. - next_s0_inflight = jax.lax.select( - is_valid & using_slot_0, - jnp.int32(1), - s0_inflight, - ) - next_s1_inflight = jax.lax.select( - is_valid & (~using_slot_0), - jnp.int32(1), - s1_inflight, - ) - return (next_s0_inflight, next_s1_inflight) - - # loop over bt, could be for loop, BT is static anyway, unroll - final_s0_inflight, final_s1_inflight = jax.lax.fori_loop( - 0, - BT, - process_decode, - (jnp.int32(0), jnp.int32(0)), - ) - - # Drain any remaining async store DMAs (at most one per slot). - @pl.when(final_s0_inflight > 0) - def _drain_slot_0(): - temp_desc = pltpu.make_async_copy( - src_ref=decode_store_scratch.at[pl.ds(0, 1)], - dst_ref=recurrent_state_out.at[pl.ds(0, 1)], - sem=decode_write_semaphore.at[0], - ) - temp_desc.wait() - - @pl.when(final_s1_inflight > 0) - def _drain_slot_1(): - temp_desc = pltpu.make_async_copy( - src_ref=decode_store_scratch.at[pl.ds(1, 1)], - dst_ref=recurrent_state_out.at[pl.ds(0, 1)], - sem=decode_write_semaphore.at[1], - ) - temp_desc.wait() - - # Mask and write accumulated outputs to HBM - mask = (jnp.arange(BT) - < decode_count).astype(decode_output_scratch.dtype)[:, None] - decode_output_scratch_masked = decode_output_scratch[...] * mask - decode_output_ref[...] = decode_output_scratch_masked - + decode_processor.process() return None # Prefill Branch # Process prefill if there is valid prefill work in this step @pl.when(prefill_valid > 0) def process_prefill(): - # TODO: eliminate k.transpose in matmuls by directly slicing in the right shape above - - # not used meaningfully, because dma is sync. - # intention is to index into scratch for storing state and not overwrite each other - prefill_slot = prefill_req_id % 2 - - def process_regular_prefill(): - init_has_init = has_initial_state[prefill_req_id][...] - init_state_idx = state_indices[prefill_req_id][...] - should_load_init = (is_first_chunk > 0) & (init_has_init > 0) - should_zero_init = (is_first_chunk > 0) & (init_has_init == 0) - - @pl.when(should_load_init) - def _start_init_load(): - copy_op = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(init_state_idx, 1)], - dst_ref=state_commit_scratch, - sem=prefill_semaphore.at[prefill_slot], - ) - copy_op.start() - - @pl.when(should_zero_init) - def _zero_init_state(): - prefill_scratch[prefill_slot] = jnp.zeros( - (n_v, d_k, d_v), dtype=prefill_scratch.dtype) - - ### Preparataion for chunk wise math, - ### this kernel design could be optimized lot by not doing this every chunk - # 1. Extract Q, K, V, g, beta for the chunk - key_dim = n_kq * d_k - - # Workaround: Upcast to fp32 to avoid NaNs in long sequences - qkv_chunk = prefill_qkv_ref[...].astype(jnp.float32) # (C, d) - # Fused SiLU - qkv_chunk = jax.nn.silu(qkv_chunk) - q = qkv_chunk[:, :key_dim] - k = qkv_chunk[:, key_dim:2 * key_dim] - v = qkv_chunk[:, 2 * key_dim:] - - # Load a, b - a_raw_chunk = prefill_a_raw_ref[...] # (C, 128) - b_raw_chunk = prefill_b_raw_ref[...] # (C, 128) - - # Slice and transpose to match expected shape (n_v, C), - # TODO: this transpose can be eliminated - a_raw_processed = a_raw_chunk[:, :n_v].T - b_raw_processed = b_raw_chunk[:, :n_v].T - - a_raw_processed = a_raw_processed.astype(jnp.float32) - b_raw_processed = b_raw_processed.astype(jnp.float32) - beta = jax.nn.sigmoid(b_raw_processed) - g = -jnp.exp(a_log_ref[...][:, None].astype( - jnp.float32)) * jax.nn.softplus(a_raw_processed + dt_bias_ref[ - ...][:, None].astype(jnp.float32)) - # Workaround: Clamp g to avoid underflow to negative inf - # g is always negative, from above line - # for long prefill sequence this negative value will get more negative - # pow(e,-100) is close to 0. - g = jnp.maximum(g, -100.0) - prefill_count = schedule_table[step, 3][...] - mask_float = (jnp.arange(C) < prefill_count).astype(q.dtype) - q = jnp.where(mask_float[:, None] > 0, q, 0.0) - k = jnp.where(mask_float[:, None] > 0, k, 0.0) - g = jnp.where(mask_float[None, :] > 0, g, 0.0) - v = jnp.where(mask_float[:, None] > 0, v, 0.0) - beta = jnp.where(mask_float[None, :] > 0, beta, 0.0) - - q = q.reshape(C, n_kq, d_k) - k = k.reshape(C, n_kq, d_k) - v = v.reshape(C, n_v, d_v) - - if use_qk_norm_in_gdn: - q = l2_normalize(q) - k = l2_normalize(k) - - # Transpose first, then repeat: the transpose operates on the - # smaller (C, n_kq, d_k) tensor; the subsequent repeat on the - # (now-leading) axis produces the same final layout. - q = q.transpose(1, 0, 2) - k = k.transpose(1, 0, 2) - v = v.transpose(1, 0, 2) - - repeat_factor = n_v // n_kq - if repeat_factor > 1: - q = jnp.repeat(q, repeat_factor, axis=0) - k = jnp.repeat(k, repeat_factor, axis=0) - - scale = d_k**-0.5 - q = q * scale - - g_cumsum_list = [] - current_sum = jnp.zeros((n_v, ), dtype=jnp.float32) - # cumsum not implemented in pallas - for i in range(C): - current_sum = current_sum + g[:, i].astype(jnp.float32) - g_cumsum_list.append(current_sum) - g_cumsum = jnp.stack(g_cumsum_list, axis=-1) - k_beta = k * beta[..., None] - - # Fuse S and S_q into a single matmul. - k_T = k.transpose(0, 2, 1) - kbeta_q = jnp.concatenate([k_beta, q], axis=1) # (n_v, 2C, d_k) - S_both = jnp.matmul( - kbeta_q.astype(jnp.float32), - k_T.astype(jnp.float32), - precision=jax.lax.Precision.HIGHEST, - ) # (n_v, 2C, C) - S = S_both[:, :C, :] - S_q = S_both[:, C:, :] - - g_diff = g_cumsum[..., :, None] - g_cumsum[..., None, :] - i = jnp.arange(C)[:, None] - j = jnp.arange(C)[None, :] - mask_float = (i > j).astype(jnp.float32) - - # Defensive code to handle large positive g_diff which can cause - # overflow in exp, - # TODO: analyze if this is a common case and if we can remove this or do - # by other means (like clipping g values before cumsum or using a - # different data type for g/g_cumsum) - g_diff_safe = jnp.minimum(g_diff, 0.0) - S = jnp.where(mask_float[None, :, :] > 0, S * jnp.exp(g_diff_safe), - 0.0) - - mask_float_q = (i >= j).astype(jnp.float32) - g_diff_Sq = g_diff_safe * mask_float_q[None, ...] + ( - 1.0 - mask_float_q[None, ...]) * (-1e30) - S_q = S_q * jnp.exp(g_diff_Sq) - S_q = S_q * mask_float_q[None, ...] - - I_plus_S = jnp.eye(C, dtype=jnp.float32)[None, ...] + S - # TODO: call the function in kernels file - A_inv = invert_triangular_matrix(I_plus_S, block_size=16) - - # Fuse u and w into a single matmul. Both compute - # A_inv @ ; stack v_beta and k_beta_g along the last - # axis (d_v -> d_v + d_k), do one matmul, then split. - v_beta = v * beta[..., None] - k_beta_g = k_beta * jnp.exp(g_cumsum)[..., None] - vk_in = jnp.concatenate( - [ - v_beta.astype(jnp.float32), - k_beta_g.astype(jnp.float32), - ], - axis=2, - ) # (n_v, C, d_v + d_k) - uw = jnp.matmul(A_inv, vk_in, precision=jax.lax.Precision.HIGHEST) - u = uw[..., :d_v] - w = uw[..., d_v:] - - q_g = q * jnp.exp(g_cumsum)[..., None] - - # Fuse attn_inter and v_prime into a single matmul. With attentionDP, the async copy leads to very small perf gain. - @pl.when(should_load_init) - def _finish_init_load(): - temp_desc = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(init_state_idx, 1)], - dst_ref=state_commit_scratch, - sem=prefill_semaphore.at[prefill_slot], - ) - temp_desc.wait() - prefill_scratch[prefill_slot] = state_commit_scratch[0].astype( - prefill_scratch.dtype) - - current_state = prefill_scratch[prefill_slot] - qw = jnp.concatenate([q_g.astype(jnp.float32), w], axis=1) - comb = jnp.matmul( - qw, - current_state.astype(jnp.float32), - precision=jax.lax.Precision.HIGHEST, - ) # (n_v, 2C, d_v) - attn_inter = comb[:, :C, :] - v_prime = comb[:, C:, :] - - v_new = u - v_prime - term2 = jnp.matmul(S_q, v_new, precision=jax.lax.Precision.HIGHEST) - o_c = attn_inter + term2 - - g_i_last_exp = jnp.exp(g_cumsum[..., -1, None, None]) - g_diff_exp_state = jnp.exp(g_cumsum[..., -1, None] - - g_cumsum)[..., None] - k_i_g_diff = k * g_diff_exp_state - - update_term = jnp.matmul( - k_i_g_diff.transpose(0, 2, 1).astype(jnp.float32), - v_new, - precision=jax.lax.Precision.HIGHEST, - ) - h_new = current_state * g_i_last_exp + update_term - - prefill_scratch[prefill_slot] = h_new.astype(prefill_scratch.dtype) - - # Store state only if it's the last chunk of the request. - store_state_idx = state_indices[prefill_req_id][...] - - @pl.when(is_last_chunk > 0) - def store_state(): - # TODO: if dtype of state in HBM is always f32, - # then we can eliminate this copy and directly write from scratch to HBM - state_commit_scratch[0] = prefill_scratch[prefill_slot].astype( - state_commit_scratch.dtype) - copy_op = pltpu.make_async_copy( - src_ref=state_commit_scratch, - dst_ref=recurrent_state_out.at[pl.ds(store_state_idx, 1)], - sem=prefill_semaphore.at[prefill_slot], - ) - copy_op.start() - copy_op.wait() - - # TODO: eliminate this transpose and reshape by directly writing in the right shape above - o_c_tr = o_c.transpose(1, 0, 2) - o_c_flat = o_c_tr.reshape(C, n_v * d_v) - - prefill_count = schedule_table[step, 3][...] - mask_float = (jnp.arange(C) < prefill_count).astype(o_c_flat.dtype) - o_c_flat_masked = o_c_flat * mask_float[:, None] - prefill_output_ref[...] = o_c_flat_masked.astype( - prefill_output_ref.dtype) - - return None - - def process_transition_prefill(): - # this is processing prefill sequences in a sublane that has multiple sequences - C_trans = sublanesize - key_dim = n_kq * d_k - - first_req_id = schedule_table[step, 11][...] - first_is_first = schedule_table[step, 11 + C_trans][...] - first_slot = first_req_id % 2 - first_has_init = has_initial_state[first_req_id][...] - should_load_first = (first_is_first > 0) & (first_has_init > 0) - first_state_idx = state_indices[first_req_id][...] - - # Async initial load - @pl.when(should_load_first) - def _start_first_load(): - copy_op = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(first_state_idx, 1)], - dst_ref=state_commit_scratch, - sem=prefill_semaphore.at[first_slot], - ) - copy_op.start() - - # Workaround: Upcast to fp32 to avoid NaNs - qkv_chunk = prefill_qkv_ref[:C_trans, :].astype(jnp.float32) - # Fused SiLU TODO: maybe 'SiLU' needs to be parametrized, - qkv_chunk = jax.nn.silu(qkv_chunk) - q = qkv_chunk[:, :key_dim] - k = qkv_chunk[:, key_dim:2 * key_dim] - v = qkv_chunk[:, 2 * key_dim:] - - # Load untransposed a and b - a_raw_chunk = prefill_a_raw_ref[...] # (C, 128) - b_raw_chunk = prefill_b_raw_ref[...] # (C, 128) - - # Slice and transpose to match expected shape (n_v, C_trans) - a_raw_processed = a_raw_chunk[:C_trans, :n_v].T - b_raw_processed = b_raw_chunk[:C_trans, :n_v].T - - a_raw_processed = a_raw_processed.astype(jnp.float32) - b_raw_processed = b_raw_processed.astype(jnp.float32) - beta_chunk = jax.nn.sigmoid(b_raw_processed) - g_chunk = -jnp.exp(a_log_ref[...][:, None].astype( - jnp.float32)) * jax.nn.softplus(a_raw_processed + dt_bias_ref[ - ...][:, None].astype(jnp.float32)) - g_chunk = jnp.maximum(g_chunk, -100.0) - q = q.reshape(C_trans, n_kq, d_k) - k = k.reshape(C_trans, n_kq, d_k) - v = v.reshape(C_trans, n_v, d_v) - - if use_qk_norm_in_gdn: - q = l2_normalize(q) - k = l2_normalize(k) - - # Transpose first, then repeat - q = q.transpose(1, 0, 2) - k = k.transpose(1, 0, 2) - v = v.transpose(1, 0, 2) - - repeat_factor = n_v // n_kq - if repeat_factor > 1: - q = jnp.repeat(q, repeat_factor, axis=0) - k = jnp.repeat(k, repeat_factor, axis=0) - - scale = d_k**-0.5 - q = q * scale - - # Cold-start: zero the slot in place when the first sequence has - # no carried-over state, then read. - @pl.when((first_is_first > 0) & (first_has_init == 0)) - def _zero_first_slot(): - prefill_scratch[first_slot] = jnp.zeros( - (n_v, d_k, d_v), dtype=prefill_scratch.dtype) - - # Finish the async initial load started above - @pl.when(should_load_first) - def _finish_first_load(): - temp_desc = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(first_state_idx, 1)], - dst_ref=state_commit_scratch, - sem=prefill_semaphore.at[first_slot], - ) - temp_desc.wait() - prefill_scratch[first_slot] = state_commit_scratch[0].astype( - prefill_scratch.dtype) - - h = prefill_scratch[first_slot] - - current_r = first_req_id - sequence_valid = True - - # Loop over tokens in the sublane. - for i in range(sublanesize): - t_req = schedule_table[step, 11 + i][...] - t_is_first = schedule_table[step, 11 + C_trans + i][...] - t_is_last = schedule_table[step, 11 + 2 * C_trans + i][...] - - is_new_seq = t_req != current_r - sequence_valid = jnp.where(is_new_seq, True, sequence_valid) - - # Ignore tokens that belong to decode requests, - # (assumes decode tokens are at packed at head) - is_decode_token = t_req < decode_tokens - sequence_valid = jnp.where(is_decode_token, False, - sequence_valid) - - c_slot = current_r % 2 - - # Commit the previous iter's h to the current request's slot. - # The other slot is untouched. - prefill_scratch[c_slot] = h - - def do_write(): - state_idx = state_indices[current_r][...] - # Stage h only when we actually DMA out. - state_commit_scratch[0] = h.astype( - state_commit_scratch.dtype) - copy_op = pltpu.make_async_copy( - src_ref=state_commit_scratch, - dst_ref=recurrent_state_out.at[pl.ds(state_idx, 1)], - sem=prefill_semaphore.at[c_slot], - ) - copy_op.start() - copy_op.wait() - return None - - is_current_r_prefill = current_r >= decode_tokens - should_write = is_current_r_prefill & is_new_seq - jax.lax.cond(should_write, do_write, lambda: None) - - t_slot = t_req % 2 - t_has_init = has_initial_state[t_req][...] - - def load_t_state(): - state_idx = state_indices[t_req][...] - copy_op = pltpu.make_async_copy( - src_ref=recurrent_state_in.at[pl.ds(state_idx, 1)], - dst_ref=state_commit_scratch, - sem=prefill_semaphore.at[t_slot], - ) - copy_op.start() - copy_op.wait() - prefill_scratch[t_slot] = state_commit_scratch[0].astype( - prefill_scratch.dtype) - - should_load_t = (t_is_first > 0) & (t_has_init > 0) - jax.lax.cond(should_load_t, load_t_state, lambda: None) - - # Cold-start: zero the slot in place for a new sequence with - # no carried-over state. Mutually exclusive with load_t_state. - @pl.when((t_is_first > 0) & (t_has_init == 0)) - def _zero_t_slot(): - prefill_scratch[t_slot] = jnp.zeros( - (n_v, d_k, d_v), dtype=prefill_scratch.dtype) - - h = prefill_scratch[t_slot] - - current_r = t_req - - k_i = k[:, i, :] - v_i = v[:, i, :] - g_i = g_chunk[:, i] - beta_i = beta_chunk[:, i] - q_i = q[:, i, :] - - decay = jnp.exp(g_i)[..., None] - - k_state = jnp.sum(k_i[..., None] * h, axis=1) - v_diff = v_i - decay * k_state - v_new = beta_i[:, None] * v_diff - - q_state = jnp.sum(q_i[..., None] * h, axis=1) - q_k = jnp.sum(q_i * k_i, axis=-1, keepdims=True) - - out_i = decay * q_state + q_k * v_new - - k_v_new = k_i[..., None] * v_new[:, None, :] - h_new = h * decay[..., None] + k_v_new - - h = jnp.where(sequence_valid, h_new, h) - - # Mask output before invalidating the sequence for the next token. - out_i = jnp.where(sequence_valid, out_i, 0.0) - - sequence_valid = jnp.where(t_is_last > 0, False, - sequence_valid) - - prefill_output_ref[i, :] = out_i.reshape(n_v * d_v).astype( - prefill_output_ref.dtype) - - final_slot = current_r % 2 - prefill_scratch[final_slot] = h - - is_current_r_prefill = current_r >= decode_tokens - - # Store state if the current request is a prefill. - # At the end of the transition prefill. No need to make this async. - @pl.when(is_current_r_prefill) - def do_final_write(): - state_idx = state_indices[current_r][...] - # Stage h into state_commit only when we actually DMA out. - state_commit_scratch[0] = h.astype(state_commit_scratch.dtype) - copy_op = pltpu.make_async_copy( - src_ref=state_commit_scratch, - dst_ref=recurrent_state_out.at[pl.ds(state_idx, 1)], - sem=prefill_semaphore.at[final_slot], - ) - copy_op.start() - copy_op.wait() - return None - - return None - - is_transition = schedule_table[step, 10][...] - - def process_prefill_dispatch(): - return jax.lax.cond( - is_transition > 0, - lambda _: process_transition_prefill(), - lambda _: process_regular_prefill(), - operand=None, - ) - - process_prefill_dispatch() + prefill_processor.process() return None - # For transition block at boundary of decode and prefill we will have overlap - # decode block BT contains prefill tokens - # sublane size transition prefill block contains some decode tokens in the sublane - # so we need to stitch the outputs so they don't overwrite each other in global index - # we exchange decode and prefill outputs so - # prefill output ref has decode token outputs at decode token indexes in its out ref - # decode output ref has prefill token outputs have prefill token indexes in its out ref + # For transition block at boundary of decode and prefill we will have + # overlap: + # - Decode block BT contains prefill tokens + # - Sublane size transition prefill block contains some decode tokens in the + # sublane + # So we need to stitch the outputs so they don't overwrite each other in the + # global index. We exchange decode and prefill outputs so: + # - Prefill output ref has decode token outputs at decode token indexes in its + # out ref + # - Decode output ref has prefill token outputs at prefill token indexes in + # its out ref def do_stitch(): local_start = prefill_offset - decode_offset local_split = decode_tokens - prefill_offset @@ -952,7 +295,7 @@ def do_stitch(): iota = jax.lax.broadcasted_iota(jnp.int32, (sublanesize, ), 0) is_decode_mask = (iota < local_split).astype(jnp.int32)[:, None] - # 4. Merge the tensors + # 4. Merge merged_overlap = jnp.where(is_decode_mask, decode_overlap, prefill_arr) decode_output_ref[ @@ -971,7 +314,6 @@ def get_qkv_index_map_v2( schedule_table, valid_col, offset_col, - count_col, alignment=16, block_size=64, sink_offset=0, @@ -1001,9 +343,8 @@ def create_block_specs( prefill_qkv_index_map = functools.partial( get_qkv_index_map_v2, schedule_table=schedule_table, - valid_col=0, - offset_col=1, - count_col=3, + valid_col=recurrent_scan_impl.COL_PREFILL_VALID, + offset_col=recurrent_scan_impl.COL_PREFILL_OFFSET, alignment=alignment, block_size=chunk_size, sink_offset=sink_offset, @@ -1012,9 +353,8 @@ def create_block_specs( decode_qkv_index_map = functools.partial( get_qkv_index_map_v2, schedule_table=schedule_table, - valid_col=4, - offset_col=5, - count_col=7, + valid_col=recurrent_scan_impl.COL_DECODE_VALID, + offset_col=recurrent_scan_impl.COL_DECODE_OFFSET, alignment=alignment, block_size=BT, sink_offset=sink_offset, @@ -1249,7 +589,7 @@ def recurrent_scan( each request in mixed_qkv. state_indices: jax.Array of shape [num_requests] or larger. Mapping from request ID to state index. - distribution: jax.Array of shape [2]. Contains [decode_tokens, + distribution: jax.Array of shape [3]. Contains [decode_tokens, total_tokens]. n_kq: Number of query/key heads. n_v: Number of value heads. @@ -1278,7 +618,7 @@ def recurrent_scan( # Default the scoped VMEM ceiling. This value could be tuned for different state cache numerics and chunk sizes. if vmem_limit_bytes is None: - vmem_limit_bytes = tpu_info.vmem_capacity_bytes + vmem_limit_bytes = int(tpu_info.vmem_capacity_bytes * 0.8) # Pad token dimension so invalid pipeline steps DMA into a safe sink area. # Sink offset must be aligned to sublanesize for Mosaic tile compatibility. From 12534565247c697115ef9c9406dbf7bb5fa0f31e Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Mon, 8 Jun 2026 10:50:06 +0000 Subject: [PATCH 28/37] [skip ci] Update nightly support matrices for vllm (v6e/v7x) Signed-off-by: Buildkite Bot --- .../nightly/v6e/vllm/parallelism_support_matrix.csv | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv b/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv index 7abe71b054..a6cc9507ed 100644 --- a/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv +++ b/support_matrices/nightly/v6e/vllm/parallelism_support_matrix.csv @@ -1,6 +1,6 @@ Feature,Single-Host CorrectnessTest,Single-Host PerformanceTest,Multi-Host CorrectnessTest,Multi-Host PerformanceTest "CP",❓ Untested,❓ Untested,❓ Untested,❓ Untested -"DP",❌ Failing,❓ Untested,❓ Untested,❓ Untested +"DP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "EP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "PP",✅ Passing,✅ Passing,✅ Passing,✅ Passing "SP",✅ Passing,❓ Untested,❓ Untested,❓ Untested From 0dd4e0ecf76d2dff191216b26f81cded1c011917 Mon Sep 17 00:00:00 2001 From: guowei-dev Date: Mon, 8 Jun 2026 08:49:00 -0700 Subject: [PATCH 29/37] Fix DeepSeek-V4 imports for vLLM DeepseekV4Attention refactor (#2829) Signed-off-by: guowei-dev --- .buildkite/vllm_lkg.version | 2 +- .../vllm/custom_ops/deepseek_v4_attention.py | 128 +++++++----------- .../models/vllm/vllm_model_wrapper.py | 4 +- tpu_inference/runner/kv_cache_manager.py | 6 +- 4 files changed, 53 insertions(+), 87 deletions(-) diff --git a/.buildkite/vllm_lkg.version b/.buildkite/vllm_lkg.version index fd0dedc5a5..b98845e9e0 100644 --- a/.buildkite/vllm_lkg.version +++ b/.buildkite/vllm_lkg.version @@ -1 +1 @@ -063ce98fb7104dba93f72d56c094fb8b708bd793 +4efd6ffde09477800294a8ed9cc752017812c3b1 diff --git a/tpu_inference/layers/vllm/custom_ops/deepseek_v4_attention.py b/tpu_inference/layers/vllm/custom_ops/deepseek_v4_attention.py index 6a016530af..d30f4cd45b 100644 --- a/tpu_inference/layers/vllm/custom_ops/deepseek_v4_attention.py +++ b/tpu_inference/layers/vllm/custom_ops/deepseek_v4_attention.py @@ -13,27 +13,27 @@ # limitations under the License. """TPU interception for DeepSeek-V4 MLA attention (torchax path). -``DeepseekV4MLA`` is a plain ``nn.Module`` that vLLM instantiates directly -(``self.mla_attn = DeepseekV4MLA(...)`` in ``deepseek_v4/amd/model.py``). Unlike -the MHC ops or the attention-impl bases, it is NOT a vLLM ``CustomOp`` and has no +``DeepseekV4Attention`` is a plain ``nn.Module`` that vLLM instantiates directly +(the AMD decoder does ``self.attn = DeepseekV4ROCMAiterMLAAttention(...)``, a +``DeepseekV4Attention`` subclass, in ``deepseek_v4/amd/model.py``). Unlike the +MHC ops or the attention-impl bases, it is NOT a vLLM ``CustomOp`` and has no ``register_oot`` hook, so there is no registry-based way to swap it. Its -constructor is also CUDA-bound (asserts a CUDA device capability and allocates -``torch.cuda.Event``), so it cannot run on TPU as-is. +constructor is also CUDA-bound (allocates ``torch.cuda.Event``), so it cannot +run on TPU as-is. Instead we substitute the class symbol before the model is built. Because -``amd/model.py`` does ``from ...attention import DeepseekV4MLA``, the name is -bound into the ``amd.model`` module namespace at import time; patching -``attention.DeepseekV4MLA`` alone would not take effect. ``patch_deepseek_v4_mla_cls`` +``amd/model.py`` does ``from ...amd.rocm import DeepseekV4ROCMAiterMLAAttention``, +the name is bound into the ``amd.model`` module namespace at import time; patching +it on ``amd.rocm`` alone would not take effect. ``patch_deepseek_v4_mla_cls`` rebinds it on ``amd.model`` directly. It is invoked from ``_maybe_patch_for_deepseek_v4`` in ``vllm_model_wrapper`` while ``is_rocm`` is forced True and the package has been reloaded onto the AMD implementation. """ import torch import torch.nn as nn -from vllm.config import CacheConfig, VllmConfig -from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.models.deepseek_v4.attention import DeepseekV4MLA +from vllm.config import VllmConfig +from vllm.model_executor.models.utils import extract_layer_index +from vllm.models.deepseek_v4.attention import DeepseekV4Attention from vllm.v1.attention.backend import AttentionBackend from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec @@ -44,20 +44,39 @@ logger = init_logger(__name__) -class VllmDeepseekV4MLAAttention(nn.Module, AttentionLayerBase): +class VllmDeepseekV4MLAAttention(DeepseekV4Attention): def __init__( self, - head_dim: int, - compress_ratio: int, - prefix: str, - cache_config: CacheConfig, + vllm_config: VllmConfig, + prefix: str = "", + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list | None = None, ) -> None: nn.Module.__init__(self) + config = vllm_config.model_config.hf_config self.prefix = prefix - self.head_dim = head_dim - self.compress_ratio = compress_ratio - self.cache_dtype = cache_config.cache_dtype + self.head_dim = config.head_dim + self.cache_dtype = vllm_config.cache_config.cache_dtype + layer_id = extract_layer_index(prefix) + if layer_id < config.num_hidden_layers: + self.compress_ratio = max(1, config.compress_ratios[layer_id]) + else: + self.compress_ratio = 1 + + # Abstract platform hooks required to instantiate the DeepseekV4Attention + # ABC; unused on the TPU pass-through path. + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + return num_heads + + def forward_mqa(self, q: torch.Tensor, kv: torch.Tensor, + positions: torch.Tensor, output: torch.Tensor) -> None: + raise NotImplementedError + + def _o_proj(self, o: torch.Tensor, + positions: torch.Tensor) -> torch.Tensor: + raise NotImplementedError def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: return MLAAttentionSpec( @@ -77,61 +96,6 @@ def process_weights_after_loading(self, act_order: bool = False) -> None: def get_attn_backend(self) -> type[AttentionBackend]: return PallasMLAttentionBackend - def forward( - self, - q: torch.Tensor, # [T, num_heads, head_dim] - kv: torch.Tensor, # [T, 1, head_dim] - positions: torch.Tensor, - output: torch.Tensor, # [T, num_heads, head_dim] - ) -> None: - logger.error( - "DeepseekV4MLA.forward is not implemented, just a pass-through for now" - ) - return q - - -class VllmDeepseekV4MLA(DeepseekV4MLA): - - def __init__( - self, - hidden_size: int, - num_heads: int, - head_dim: int, - scale: float, - qk_nope_head_dim: int, - qk_rope_head_dim: int, - v_head_dim: int, - q_lora_rank: int | None, - kv_lora_rank: int, - o_lora_rank: int | None, - vllm_config: VllmConfig, - fused_wqa_wkv: torch.nn.Module, - q_norm: torch.nn.Module, - wq_b: torch.nn.Module, - kv_norm: torch.nn.Module, - wo_a: torch.nn.Module, - wo_b: torch.nn.Module, - attn_sink: torch.nn.Module, - rotary_emb: torch.nn.Module, - indexer: torch.nn.Module | None, - indexer_rotary_emb: torch.nn.Module, - topk_indices_buffer: torch.Tensor | None, - aux_stream_list: list | None, - window_size: int, - compress_ratio: int | None, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - nn.Module.__init__(self) - - self.mla_attn = VllmDeepseekV4MLAAttention( - head_dim=head_dim, - compress_ratio=compress_ratio, - prefix=prefix, - cache_config=cache_config, - ) - def forward( self, positions: torch.Tensor, @@ -139,18 +103,20 @@ def forward( llama_4_scaling: torch.Tensor | None = None, ) -> torch.Tensor: logger.error( - "VllmDeepseekV4MLA.forward is not implemented, just a pass-through for now" + "VllmDeepseekV4MLAAttention.forward is not implemented, just a pass-through for now" ) return hidden_states def patch_deepseek_v4_mla_cls() -> None: - """Rebind ``DeepseekV4MLA`` to the TPU subclass for DS V4 model module. + """Rebind ``DeepseekV4ROCMAiterMLAAttention`` to the TPU subclass. Must run after ``vllm.models.deepseek_v4.amd.model`` is imported (it holds - its own ``from ...attention import DeepseekV4MLA`` reference) and before the - model is constructed. + its own ``from ...amd.rocm import DeepseekV4ROCMAiterMLAAttention`` + reference) and before the model is constructed. """ import vllm.models.deepseek_v4.amd.model as ds_v4_amd_model - ds_v4_amd_model.DeepseekV4MLA = VllmDeepseekV4MLA - logger.info("Patched DeepseekV4MLA -> VllmDeepseekV4MLA for TPU.") + ds_v4_amd_model.DeepseekV4ROCMAiterMLAAttention = VllmDeepseekV4MLAAttention + logger.info( + "Patched DeepseekV4ROCMAiterMLAAttention -> VllmDeepseekV4MLAAttention for TPU." + ) diff --git a/tpu_inference/models/vllm/vllm_model_wrapper.py b/tpu_inference/models/vllm/vllm_model_wrapper.py index e5c7b6bc01..bd65507939 100644 --- a/tpu_inference/models/vllm/vllm_model_wrapper.py +++ b/tpu_inference/models/vllm/vllm_model_wrapper.py @@ -116,8 +116,8 @@ def _maybe_patch_for_deepseek_v4(vllm_config: VllmConfig): _ml_utils._MODEL_ARCH_BY_HASH.clear() _try_load_model_cls.cache_clear() - # DeepseekV4MLA is a plain nn.Module instantiated directly in - # amd/model.py (not CustomOp/register_oot hook). + # DeepseekV4ROCMAiterMLAAttention is a plain nn.Module instantiated directly + # in amd/model.py (not CustomOp/register_oot hook). # Swap the class symbol for the TPU subclass before the # model is built. from tpu_inference.layers.vllm.custom_ops.deepseek_v4_attention import \ diff --git a/tpu_inference/runner/kv_cache_manager.py b/tpu_inference/runner/kv_cache_manager.py index d34f0ac50c..23aa1e988b 100644 --- a/tpu_inference/runner/kv_cache_manager.py +++ b/tpu_inference/runner/kv_cache_manager.py @@ -24,8 +24,8 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.mamba.abstract import MambaBase from vllm.model_executor.layers.mla import MLAAttention -from vllm.models.deepseek_v4.attention import (DeepseekV4IndexerCache, - DeepseekV4MLAAttention) +from vllm.models.deepseek_v4.attention import (DeepseekV4Attention, + DeepseekV4IndexerCache) from vllm.models.deepseek_v4.compressor import CompressorStateCache from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache @@ -61,7 +61,7 @@ def is_cache_for_ds_v4(attn_module: AttentionLayerBase) -> bool: return isinstance(attn_module, DeepseekV4IndexerCache) or isinstance( attn_module, DeepseekV4SWACache) or isinstance( - attn_module, DeepseekV4MLAAttention) or isinstance( + attn_module, DeepseekV4Attention) or isinstance( attn_module, CompressorStateCache) From a342edc38cfd5d108ffec3dd3693af3f0dd5da7d Mon Sep 17 00:00:00 2001 From: Juncheng Gu Date: Mon, 8 Jun 2026 22:59:31 -0700 Subject: [PATCH 30/37] [Feat][Disagg] Use tpu-raiden for KV cache transfer in disaggregated serving (#2837) Signed-off-by: Juncheng Gu --- .../distributed/tpu_raiden_connector.py | 637 ++++++++++++++++++ 1 file changed, 637 insertions(+) create mode 100644 tpu_inference/distributed/tpu_raiden_connector.py diff --git a/tpu_inference/distributed/tpu_raiden_connector.py b/tpu_inference/distributed/tpu_raiden_connector.py new file mode 100644 index 0000000000..e5b9634bc7 --- /dev/null +++ b/tpu_inference/distributed/tpu_raiden_connector.py @@ -0,0 +1,637 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Proxy server routes the request to P with max_output_tokens=1 + +P workflow: + P recives the request + + P scheduler checks if the prefill is full done in `request_finished()` + If done: + P puts the request-id in `scheduler_output.finished_req_ids` + and puts the request in `scheduler_output.kv_connector_metadata.reqs_to_send` + P responds the proxy server with `finished_req_ids` and the `kv_transfer_params` + P worker gets `reqs_to_send` and runs async `_prepare_kv_and_wait()` + Else: + P schedules the prefill with multiple turns due to chunked-prefill. + + P worker checks if the request has been pulled by D + If done: + P worker puts the request-id in `done_sending()` + P scheduler frees blocks for the requet in done sending. + Else: + P holds the blocks for the request until it's pulled by D + + ( + One scheduler step can finish: + scheduler RUNNING -> connector reqs_to_send -> worker prefill -> output + The waiting buffer will get freed after notified by D or expired. + ) + +Proxy server recives the response from P and forwards it to D + +D workflow: + D recives the request + + D scheduler calculates the num of tokens needing to pull from P in `get_num_new_matched_tokens()` + D checks if need to pull from P + If true: + D puts the request in `scheduler_output.kv_connector_metadata.reqs_to_load` + D worker gets `reqs_to_load` and runs `_pull_and_write_kv()` in separate threads (to be async) + D worker checks if the async loading is done: + If done: + D worker puts the request-id in `done_recving`. + D scheduler then knows the request can be scheduled for decoding now. The model decode + will happen in the next scheduler step. + Else: + D worker handles other requests first. + Else (too short prompt, full local prefix-cache): + D still needs to puts the request in `reqs_to_load` but with None metadata, because D needs to + notify P the prefilled KV cache is no longer needed and can be freed in P. + + ( + Two scheduler steps can finish: + scheduler WAITING_FOR_REMOTE_KVS -> connector reqs_to_load -> worker wait for pulling + worker pulling done, notify P to free blocks + scheduler RUNNING -> connector reqs_to_load=None -> worker decode -> output + The waiting buffer will get freed after notified by D or expired. + ) +""" + +import os +import time +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Optional +from uuid import uuid4 + +import jax +from jax.sharding import Mesh +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole) +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorPromMetrics, KVConnectorStats, PromMetric, PromMetricT) +from vllm.utils.math_utils import round_down +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.request import RequestStatus + +if TYPE_CHECKING: + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.request import Request + +try: + from api.jax.transfer_engine import TransferEngine as RaidenTransferEngine + _RAIDEN_IMPORT_ERROR = None +except Exception as _exc: # pylint: disable=broad-except + RaidenTransferEngine = None + _RAIDEN_IMPORT_ERROR = _exc + +import tpu_inference.distributed.utils as dist_utils +from tpu_inference import envs +from tpu_inference.distributed.tpu_connector_stats import ( + TpuKVConnectorPromMetrics, TpuKVConnectorStats) +from tpu_inference.logger import init_logger +from tpu_inference.runner.tpu_runner import TPUModelRunner + +ReqId = str + +# Feature requests: +# 1. support async pulling natively +# 2. partial pulling (like RDMA) +# 3. non-blocking jax array read/write + +logger = init_logger(__name__) + + +@dataclass +class SendMeta: + uuid: int + # `list[int]` used for non-HMA connector + # `list[list[int]]` used for HMA connector (per-kv-cache-group) + local_block_ids: list[int] | list[list[int]] + expiration_time: float + + +@dataclass +class LoadMeta: + uuid: int + # `list[int]` used for non-HMA connector. + # `list[list[int]]` used for HMA connector (per-kv-cache-group). + local_block_ids: list[int] | list[list[int]] | None + remote_block_ids: list[int] | list[list[int]] | None + remote_host: str | list[str] + remote_port: int | list[int] + + +# The metadata used for communicating between scheduler and worker connectors. +@dataclass +class TPUConnectorMetadata(KVConnectorMetadata): + reqs_to_send: dict[ReqId, SendMeta] = field(default_factory=dict) + reqs_to_load: dict[ReqId, LoadMeta] = field(default_factory=dict) + + +class TPUConnector(KVConnectorBase_V1): + + def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole, + kv_cache_config: KVCacheConfig): + super().__init__(vllm_config, role, kv_cache_config) + assert vllm_config.kv_transfer_config is not None + self._connector_metadata = None + + if role == KVConnectorRole.SCHEDULER: + self.connector_scheduler = \ + TPUConnectorScheduler(vllm_config) + self.connector_worker = None + elif role == KVConnectorRole.WORKER: + self.connector_scheduler = None + self.connector_worker = TPUConnectorWorker(vllm_config) + + ############################################################ + # Scheduler Side Methods + ############################################################ + def get_num_new_matched_tokens( + self, request: "Request", + num_computed_tokens: int) -> tuple[int, bool]: + assert self.connector_scheduler is not None + return self.connector_scheduler.get_num_new_matched_tokens( + request, num_computed_tokens) + + def update_state_after_alloc(self, request: "Request", + blocks: "KVCacheBlocks", + num_external_tokens: int): + assert self.connector_scheduler is not None + return self.connector_scheduler.update_state_after_alloc( + request, blocks, num_external_tokens) + + def build_connector_meta( + self, + scheduler_output: SchedulerOutput, + ) -> TPUConnectorMetadata: + assert self.connector_scheduler is not None + return self.connector_scheduler.build_connector_meta() + + def request_finished( + self, + request: "Request", + block_ids: list[int], + ) -> tuple[bool, Optional[dict[str, Any]]]: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, block_ids) + + def get_finished_count(self) -> int: + assert self.connector_scheduler is not None + return self.connector_scheduler.get_finished_count() + + ############################################################ + # Worker Side Methods + ############################################################ + def register_kv_caches(self, kv_caches: list[jax.Array]): + """ + We don't register kv_caches in connector, we call `register_runner` and + use runner.kv_caches directly instead because the ref of runner.kv_caches + would be reassigned during model forward. + """ + pass + + def register_runner(self, runner: TPUModelRunner) -> None: + assert self.connector_worker is not None + self.connector_worker.register_runner(runner) + + def start_load_kv(self, _, **kwargs) -> None: + assert self.connector_worker is not None + assert isinstance(self._connector_metadata, TPUConnectorMetadata) + self.connector_worker.process_send_load(self._connector_metadata) + + def wait_for_layer_load(self, layer_name: str) -> None: + """TPU connector doesn't support layer wise load.""" + pass + + def save_kv_layer(self, *args, **kwargs) -> None: + """TPU connector doesn't support layer wise save.""" + pass + + def wait_for_save(self): + """ + Not useful for TPU, because by the design of vLLM KVConnectorModelRunnerMixin, + this function is only called when scheduler_output.total_num_scheduled_tokens is not 0. + But the reqs_to_send is only available after the req finished prefilling where the + total_num_scheduled_tokens could be 0 if no other running reqs. + So we run saving logic in `start_load_kv -> process_send_load` instead. + """ + pass + + def get_finished(self, + finished_req_ids: set[str]) -> tuple[set[str], set[str]]: + assert self.connector_worker is not None + return self.connector_worker.get_finished() + + def get_kv_connector_stats(self) -> KVConnectorStats | None: + """ + Get the KV transfer stats for the connector. + """ + if self.connector_worker is None: + return None + return self.connector_worker.get_kv_connector_stats() + + @classmethod + def build_kv_connector_stats( + cls, + data: dict[str, Any] | None = None) -> KVConnectorStats | None: + return (TpuKVConnectorStats( + data=data) if data is not None else TpuKVConnectorStats()) + + @classmethod + def build_prom_metrics( + cls, + vllm_config: VllmConfig, + metric_types: dict[type[PromMetric], type[PromMetricT]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ) -> KVConnectorPromMetrics: + return TpuKVConnectorPromMetrics(vllm_config, metric_types, labelnames, + per_engine_labelvalues) + + +class TPUConnectorScheduler(): + + def __init__(self, vllm_config: "VllmConfig"): + self.vllm_config = vllm_config + self.config = vllm_config.kv_transfer_config + self.is_producer = self.config.is_kv_producer + + self.block_size = vllm_config.cache_config.block_size + + # This is updated in self.update_state_after_alloc() for D, + # each request that needs to pull KV cache from remote will be added to it. + self.reqs_to_send: dict[ReqId, SendMeta] = {} + + # This is updated in self.request_finished() for P, + # each request that finished prefilling will be added to it. + self.reqs_to_load: dict[ReqId, LoadMeta] = {} + + self.kv_ip = dist_utils.get_kv_ips() + self.kv_port = dist_utils.get_kv_ports() + logger.info( + f"TPUConnectorScheduler --> kv_ip={self.kv_ip} | kv_port={self.kv_port}" + ) + + def get_num_new_matched_tokens( + self, + request: "Request", + num_computed_tokens: int, + ) -> tuple[int, bool]: + """ + D workers use this to get the number of new tokens + that can be loaded from remote P workers. + No-op for P workers. + + Args: + request (Request): the request object. + num_computed_tokens (int): the number of locally + computed tokens for this request + + Returns: + A tuple with the following elements: + - The number of tokens that will be loaded from the + external KV cache. + - If async loading. Must be 'False' for TPU connector + because TPU pulls KV cache in a blocking way. + + """ + if self.is_producer or not request.kv_transfer_params: + return 0, False + + # Only trigger 1 KV transfer per request. + if request.kv_transfer_params.get("do_remote_prefill", True) is False: + # logger.debug(f"TPUConnector Scheduler skip kv transfer for request {request.request_id} as it already pulled before.") + return 0, False + + assert num_computed_tokens % self.block_size == 0 + # This rounding logic must be consistent with calculating + # remote_block_ids in P's request_finished() + rounded_num_prompt_tokens = round_down(len(request.prompt_token_ids), + self.block_size) + count = max(rounded_num_prompt_tokens - num_computed_tokens, 0) + # NOTE(xiang): Although the JAX P2P pulling is a blocking op, we will run it in a + # separte thread to make it async, so we are safe to return True here. + if count > 0: + return count, True + return 0, False + + def update_state_after_alloc(self, request: "Request", + blocks: "KVCacheBlocks", + num_external_tokens: int): + """ + Update states after block allocation. + No-op for P workers. + + Args: + request (Request): the request object. + blocks (KVCacheBlocks): the blocks allocated for the request. + num_external_tokens (int): the number of tokens that will be + loaded from the external KV cache. + """ + if self.is_producer or not request.kv_transfer_params: + return + + params = request.kv_transfer_params + if num_external_tokens > 0: + # We need to load KV-cache from remote (partial prefix cache hit). + local_block_ids = blocks.get_block_ids()[0] + + # NOTE(xiang): D needs to pull the whole prefill blocks from the remote + # regardless how much ratio the prefix cache hits. + # The reason is JAX P2P doesn't work as RDMA, instead it works like: + # P just prepares the whole prefilled data and waits for pulling, then D pulls the + # whole data. Which means even with partial prefix cache hit on D, D cannot only + # pull the remaining partial data from P. + # Unless we implement a side channel to let P know the prefix cache hit info on D, + # so P can prepare those non-hit KV only, with that we need to change to: + # local_block_ids = blocks.get_unhashed_block_ids() + + self.reqs_to_load[request.request_id] = LoadMeta( + uuid=params["uuid"], + local_block_ids=local_block_ids, + remote_block_ids=params["remote_block_ids"], + remote_host=params["remote_host"], + remote_port=params["remote_port"], + ) + else: + # This branch means two cases: + # 1. We don't need to load KV-cache from remote because of full local cache. + # 2. The async pulling is done. + # In both cases we need to send notification to let P free memory. + self.reqs_to_load[request.request_id] = LoadMeta( + uuid=params["uuid"], + local_block_ids=blocks.get_block_ids()[0], + remote_block_ids=None, + remote_host=params["remote_host"], + remote_port=params["remote_port"], + ) + + # Only trigger 1 KV transfer per request. + params["do_remote_prefill"] = False + + logger.info( + f"TPUConnector Scheduler update_state_after_alloc --> reqs_to_load={self.reqs_to_load}" + ) + + def build_connector_meta(self) -> TPUConnectorMetadata: + """ + Build the scheduler metadata and pass to the downstream worker. + + This function should NOT modify fields in the scheduler_output. + Also, calling this function will reset the state of the connector. + """ + meta = TPUConnectorMetadata() + + if self.is_producer: + meta.reqs_to_send = self.reqs_to_send + self.reqs_to_send = {} + else: + meta.reqs_to_load = self.reqs_to_load + self.reqs_to_load = {} + + return meta + + def get_finished_count(self) -> int: + """ + Return how many workers need pull the kv cache and report back. + """ + return len(self.kv_ip) if isinstance(self.kv_ip, list) else 1 + + def request_finished( + self, + request: "Request", + block_ids: list[int], + ) -> tuple[bool, Optional[dict[str, Any]]]: + """ + Called when a request has finished, before its blocks are freed. + No-op for D workers. + + Args: + request (Request): the request object. + block_ids: The block IDs allocated for this request and need to be freed. + Returns: + True if the request is being saved/sent asynchronously and blocks + should not be freed until the request_id is returned from + get_finished(). + Optional KVTransferParams to be included in the request outputs + returned by the engine. + """ + if not self.is_producer: + return False, None + + # Mark the request finished only if the prefill is done and generates 1 output token. + # The request's max_tokens has been reset to 1, so it must be finished by length capped. + if request.status != RequestStatus.FINISHED_LENGTH_CAPPED: + return False, None + + # NOTE(xiang): Get computed blocks rounded by block_size. + # This indication means for the last partially filled block, we won't bother transfering + # KV-cache, will just let D run prefill locally. + all_full = request.num_computed_tokens % self.block_size == 0 + computed_block_ids = block_ids if all_full else block_ids[:-1] + + # If prompt < block_size, no transfer so free blocks immediately. + delay_free_blocks = len(computed_block_ids) > 0 + if delay_free_blocks: + uuid = get_uuid() + expiration_time = time.perf_counter( + ) + dist_utils.get_p2p_wait_pull_timeout() + self.reqs_to_send[request.request_id] = SendMeta( + uuid=uuid, + local_block_ids=computed_block_ids, + expiration_time=expiration_time) + kv_transfer_params = dict(uuid=uuid, + remote_block_ids=computed_block_ids, + remote_host=self.kv_ip, + remote_port=self.kv_port) + logger.info( + f"TPUConnector Scheduler ----> generated reqs_to_send={self.reqs_to_send} | " + f"kv_transfer_params={kv_transfer_params}") + else: + kv_transfer_params = {} + + return delay_free_blocks, kv_transfer_params + + +class TPUConnectorWorker: + + def __init__(self, vllm_config: VllmConfig): + self.vllm_config = vllm_config + self.config = vllm_config.kv_transfer_config + self.is_producer = self.config.is_kv_producer + + self.runner: TPUModelRunner = None + self.mesh: Mesh = None + self.multi_host = envs.TPU_MULTIHOST_BACKEND == "ray" + # default value for none distributed scenario + # when the topology is initialized, runner will update it + # based on topology_order_id + self.node_id = 0 + + # The Raiden transfer engine, constructed in register_runner() once the + # runner's kv_caches exist. Replaces the jax.experimental.transfer + # server + the ZMQ side channel + HostKVPool host staging. + self.engine = None + # Consumer-side: req_ids for which a real pull (submit_load) was issued, + # so the scheduler's later remote_block_ids=None notify step is a no-op. + self._submitted: set[ReqId] = set() + + self.host_ip = dist_utils.get_host_ip() + # Bind the engine control socket to the same port the scheduler + # advertises as remote_port (TPU_KV_TRANSFER_PORT). + self.kv_transfer_port = int(dist_utils.get_kv_transfer_port()) + + self.transfer_stats = TpuKVConnectorStats() + + logger.info(f"TPUConnector Worker --> init | " + f"is_producer={self.is_producer} | ip={self.host_ip} | " + f"kv_transfer_port={self.kv_transfer_port}") + + def register_runner(self, runner: TPUModelRunner): + if RaidenTransferEngine is None: + raise ImportError( + "RaidenTransferEngine is not importable. Add the tpu-raiden " + "export to PYTHONPATH so 'api.jax.transfer_engine' resolves " + "(and set RAIDEN_PRELOAD_ENGINE=1 so sitecustomize.py preloads " + f"the engine .so first). Original error: {_RAIDEN_IMPORT_ERROR}" + ) + self.node_id = runner.topology_order_id + self.runner = runner + self.mesh = runner.mesh + + kv_caches = runner.kv_caches + self.num_layers = len(kv_caches) + self.sharding = kv_caches[0].sharding + block_size = self.vllm_config.cache_config.block_size + max_blocks = self.vllm_config.model_config.max_model_len // block_size + num_slots = int(os.getenv("RAIDEN_NUM_SLOTS", "16")) + # H2H transport sockets per transfer (1 = single socket). Higher values + # parallelize the host-to-host pull to use more network bandwidth. + parallelism = int(os.getenv("RAIDEN_TRANSPORT_PARALLELISM", "1")) + # In the new tpu-raiden engine API, parallelism is a per-pull argument + # to start_read() rather than a constructor arg; stash it here. + self._parallelism = parallelism + skip_lock = os.getenv("RAIDEN_UNSAFE_SKIP_BUFFER_LOCK", + "true").lower() in ("1", "true", "yes") + + # The engine holds the physical KV buffers. The model forward updates + # them in place (donation), so the engine always serves/writes the live + # KV without re-registration (see plan, Blocker B). + self.engine = RaidenTransferEngine( + kv_caches=kv_caches, + local_control_port=self.kv_transfer_port, + max_blocks=max_blocks, + num_slots=num_slots, + timeout_s=float(dist_utils.get_p2p_wait_pull_timeout()), + unsafe_skip_buffer_lock=skip_lock, + ) + logger.info( + f"TPUConnector Worker {self.node_id} --> Raiden engine ready | " + f"ip={self.host_ip} | " + f"control_port={getattr(self.engine, 'local_control_port', None)} | " + f"data_port={getattr(self.engine, 'local_data_port', None)} | " + f"max_blocks={max_blocks} | num_slots={num_slots} | " + f"parallelism={parallelism}") + + def _remote_endpoint(self, req_meta: "LoadMeta") -> str: + host = req_meta.remote_host + port = req_meta.remote_port + if isinstance(host, list): + assert isinstance(port, list) and len(host) == len(port) + return f"{host[self.node_id]}:{port[self.node_id]}" + return f"{host}:{port}" + + def process_send_load(self, metadata: TPUConnectorMetadata): + """ + This is called in runner before calling model forward, + whenever the scheduler_output.total_num_scheduled_tokens is empty or not. + """ + reqs = metadata.reqs_to_send + if reqs: + assert self.is_producer + logger.info( + f"TPUConnector Worker {self.node_id} --> reqs_to_send={reqs}") + for req_id, req_meta in reqs.items(): + # Producer: register the prefilled blocks so D can pull them. + # Replaces select_from_kv_caches + kv_transfer_server.await_pull. + self.engine.notify_for_read(req_id, req_meta.uuid, + req_meta.local_block_ids) + + reqs = metadata.reqs_to_load + if reqs: + assert not self.is_producer + logger.info( + f"TPUConnector Worker {self.node_id} --> reqs_to_load={reqs}") + for req_id, req_meta in reqs.items(): + remote_endpoint = self._remote_endpoint(req_meta) + if req_meta.remote_block_ids is not None: + # Consumer: pull remote_block_ids straight into the local KV + # cache at local_block_ids. The engine does the H2H pull + H2D + # write directly into kv_caches -- no separate insert_kv_chunks. + # Replaces kv_transfer_server.connect + conn.pull + insert. + if req_id in self._submitted: + # Pre-allocated blocks may be re-issued; submit only once. + continue + self._submitted.add(req_id) + self.engine.start_read( + req_id=req_id, + uuid=req_meta.uuid, + remote_endpoint=remote_endpoint, + remote_block_ids=req_meta.remote_block_ids, + local_block_ids=req_meta.local_block_ids, + parallelism=self._parallelism, + ) + else: + # remote_block_ids is None => the async pull already finished + # (the engine wrote KV into local_block_ids and acked P during + # submit_load) or there was no pull (full local prefix cache). + # Nothing to do here: the producer is freed by the pull's own + # ack, or by timeout if no pull happened. Do NOT issue a 0-block + # submit_load -- the producer rejects a 0-block pull stream. + self._submitted.discard(req_id) + + def get_kv_connector_stats(self) -> KVConnectorStats | None: + """ + Get the KV transfer stats for the worker. + """ + # Clear stats for next iteration + if not self.transfer_stats.is_empty(): + return self.transfer_stats.clone_and_reset() + return None + + def get_finished(self) -> tuple[set[str], set[str]]: + # The engine's control plane reports producer completion (done_sending, + # after D acks) and consumer completion (done_recving, after H2H+H2D). + # Replaces the reqs_wait_pull/reqs_pulling bookkeeping + ZMQ side channel. + if self.engine is None: + return set(), set() + done_sending, done_recving, failed_recving = self.engine.complete_read( + ) + if failed_recving: + # Do NOT report failed receives as done_recving: vllm would then try + # to decode with KV that was never written and hit an AssertionError + # that kills the EngineCore (taking down all other requests). Leave + # them pending; the request times out at the API layer instead. + logger.error( + f"TPUConnector Worker {self.node_id} --> failed_recving={failed_recving}" + ) + if done_sending: + logger.info( + f"TPUConnector Worker {self.node_id} --> done_sending={done_sending}" + ) + if done_recving: + logger.info( + f"TPUConnector Worker {self.node_id} --> done_recving={done_recving}" + ) + return set(done_sending), set(done_recving) + + +def get_uuid() -> int: + int128 = uuid4().int + # Must be less than 64-bit int, otherwise vllm output encoder would raise error. + # use 50 bit to avoid GO trunk the int when doing JSon serialization + return int128 >> 78 From 810f8926f4a26be3b34c5ac6af81789959ea3b74 Mon Sep 17 00:00:00 2001 From: Orti Bazar Date: Mon, 8 Jun 2026 23:04:27 -0700 Subject: [PATCH 31/37] Enable exporting benchmarking data to MLCompass (#2787) --- .../cases/dev/mlcompass-example.json | 62 +++++++++ .../benchmark/scripts/benchmark_bootstrap.sh | 2 +- .../benchmark/scripts/generate_bk_pipeline.py | 3 + .../benchmark/scripts/mlcompass_export.py | 127 ++++++++++++++++++ .buildkite/benchmark/scripts/report_result.sh | 7 + .buildkite/benchmark/scripts/run_job.sh | 7 + requirements_benchmarking.txt | 1 + 7 files changed, 208 insertions(+), 1 deletion(-) create mode 100644 .buildkite/benchmark/cases/dev/mlcompass-example.json create mode 100644 .buildkite/benchmark/scripts/mlcompass_export.py diff --git a/.buildkite/benchmark/cases/dev/mlcompass-example.json b/.buildkite/benchmark/cases/dev/mlcompass-example.json new file mode 100644 index 0000000000..7a17b55d39 --- /dev/null +++ b/.buildkite/benchmark/cases/dev/mlcompass-example.json @@ -0,0 +1,62 @@ +{ + "global_env": { + "GCP_PROJECT_ID": "cloud-tpu-inference-test", + "GCS_BUCKET": "vllm-cb-storage2", + "GCP_INSTANCE_ID": "vllm-bm-inst", + "GCP_DATABASE_ID": "vllm-bm-bk-runs", + "MLCOMPASS_EXPORT_ENABLED": "true", + "MODELTAG": "NEW", + "EXPECTED_ETEL": 10800000 + }, + "benchmark_cases": [ + { + "case_name": "Qwen3-32B-dataset_random-inlen_1024-outlen_1024-max-concurrency-64", + "ci_queue": [ + "tpu_v7x_2_queue" + ], + "env": { + "MODELTAG": "NEW", + "EXPECTED_ETEL": 10800000, + "INPUT_LEN": 1024, + "OUTPUT_LEN": 1024, + "PREFIX_LEN": 0 + }, + "server_command_options": { + "command_type": "vllm_serve", + "env": { + "VLLM_USE_V1": "1", + "MODEL_IMPL_TYPE": "vllm" + }, + "args": { + "model": "Qwen/Qwen3-32B", + "max-num-seqs": 256, + "max-num-batched-tokens": 4096, + "tensor-parallel-size": { + "v7x-2": 2 + }, + "max-model-len": 2048, + "kv-cache-dtype": "fp8", + "gpu-memory-utilization": 0.98, + "additional-config": "{\"quantization\": { \"qwix\": { \"rules\": [{ \"module_path\": \".*\", \"weight_qtype\": \"float8_e4m3fn\", \"act_qtype\": \"float8_e4m3fn\"}]}}}", + "no-enable-prefix-caching": true, + "async-scheduling": true + } + }, + "client_command_options": { + "command_type": "vllm_bench_serve", + "args": { + "model": "Qwen/Qwen3-32B", + "backend": "vllm", + "dataset-name": "random", + "num-prompts": 320, + "random-input-len": 1024, + "random-output-len": 1024, + "max-concurrency": 64, + "request-rate": "inf", + "percentile-metrics": "ttft,tpot,itl,e2el", + "ignore-eos": true + } + } + } + ] +} \ No newline at end of file diff --git a/.buildkite/benchmark/scripts/benchmark_bootstrap.sh b/.buildkite/benchmark/scripts/benchmark_bootstrap.sh index 0fa560d92d..1dad82fb7f 100755 --- a/.buildkite/benchmark/scripts/benchmark_bootstrap.sh +++ b/.buildkite/benchmark/scripts/benchmark_bootstrap.sh @@ -47,7 +47,7 @@ buildkite-agent meta-data set "JOB_REFERENCE" "${JOB_REFERENCE}" upload_benchmark_pipeline() { local target_case_type="$BM_CASE_TYPE" - VLLM_COMMIT_HASH=$(get_vllm_commit_hash) + VLLM_COMMIT_HASH=${VLLM_COMMIT_HASH:-$(get_vllm_commit_hash)} buildkite-agent meta-data set "VLLM_COMMIT_HASH" "${VLLM_COMMIT_HASH}" TPU_COMMIT_HASH=$(git rev-parse HEAD) CODE_HASH="${VLLM_COMMIT_HASH}-${TPU_COMMIT_HASH}-" diff --git a/.buildkite/benchmark/scripts/generate_bk_pipeline.py b/.buildkite/benchmark/scripts/generate_bk_pipeline.py index 35ad53b603..1cb3860c2c 100644 --- a/.buildkite/benchmark/scripts/generate_bk_pipeline.py +++ b/.buildkite/benchmark/scripts/generate_bk_pipeline.py @@ -223,11 +223,14 @@ def create_benchmark_steps( # Include parent_dir in label for uniqueness step_label = f"[{parent_dir}] {agent} {file_basename}" case_parameter = f"{file_path}" + benchmark_name = file_basename else: step_env["TARGET_CASE_NAME"] = case_name # Include parent_dir in label for uniqueness step_label = f"[{parent_dir}] {agent} {file_basename} {case_name}" case_parameter = f"{file_path} {case_name}" + benchmark_name = case_name + step_env["MLCOMPASS_TEST_NAME"] = f"vllm:{agent}:{benchmark_name}" # Define step key and check for internal collisions step_safe_key = clean_key_string(step_label) diff --git a/.buildkite/benchmark/scripts/mlcompass_export.py b/.buildkite/benchmark/scripts/mlcompass_export.py new file mode 100644 index 0000000000..2ad941f844 --- /dev/null +++ b/.buildkite/benchmark/scripts/mlcompass_export.py @@ -0,0 +1,127 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import json +import os +import re +import time +import uuid + +from google.cloud import bigquery + +_GCP_PROJECT = os.getenv('MLCOMPASS_GCP_PROJECT', + 'google.com:ml-compass-benchmarks') +_BIGQUERY_DATASET = os.getenv('MLCOMPASS_BIGQUERY_DATASET', + 'benchmarks_dataset') +_BIGQUERY_TABLE = os.getenv('MLCOMPASS_BIGQUERY_TABLE', 'benchmarks_dev') + + +def get_env_commit_map() -> dict[str, dict[str, str]]: + code_hash = os.getenv('CODE_HASH') + match = re.match('^([0-9a-fA-F]+)-([0-9a-fA-F]+)-$', code_hash) + if not match: + raise RuntimeError(f'CODE_HASH value error: "{code_hash}"') + return { + 'vllm': { + 'repo': 'vllm-project/vllm', + 'commit': match[1], + 'branch': 'main', + }, + 'tpu-inference': { + 'repo': 'vllm-project/tpu-inference', + 'commit': match[2], + 'branch': 'main', + }, + } + + +def read_metrics_file(file_path: str) -> dict[str, float]: + metrics_dict: dict[str, float] = {} + with open(file_path, 'r') as file: + for line in file: + # Strip whitespace from the beginning and end of the line + line = line.strip() + # Skip empty lines and ensure the line contains '=' + if line and '=' in line: + key, value = line.split('=', 1) + # Strip extra spaces from key/value and convert value to float + metrics_dict[key.strip()] = float(value.strip()) + return metrics_dict + + +def get_links() -> dict[str, str]: + result = {} + build_number = os.getenv('BUILDKITE_BUILD_NUMBER') + job_id = os.getenv('BUILDKITE_JOB_ID') + if build_number and job_id: + result[ + 'buildkite'] = f'https://buildkite.com/organizations/tpu-commons/pipelines/tpu-inference-benchmark/builds/{build_number}/jobs/{job_id}/log' + return result + + +def export(metrics: dict[str, float]) -> None: + row_id = uuid.uuid4().hex + + test_name = os.getenv('MLCOMPASS_TEST_NAME') + if not test_name: + raise RuntimeError('MLCOMPASS_TEST_NAME env var is not found.') + + env_commit_map = get_env_commit_map() + client_info = { + 'github_commit': env_commit_map['tpu-inference']['commit'], + 'commit_branch_name': env_commit_map['tpu-inference']['branch'], + } + mlcompass_tracking_id = os.getenv('MLCOMPASS_TRACKING_ID', row_id) + mlcompass_execution_mode = os.getenv('MLCOMPASS_EXECUTION_MODE', 'oneshot') + + row = { + 'entry_id': row_id, + 'sponge_id': os.getenv('MLCOMPASS_SPONGE_ID'), + 'test_name': test_name, + 'succeeded': bool(metrics), + 'metrics': json.dumps(metrics), + 'link_map': json.dumps(get_links()), + 'exc_timestamp_millis': int(time.time() * 1000), + 'client_info': client_info, + 'env_commit_map': json.dumps(env_commit_map), + 'mlcompass_tracking_id': mlcompass_tracking_id, + 'mlcompass_execution_mode': mlcompass_execution_mode, + } + print(f'Exporting row: {row}') + + client = bigquery.Client(project=_GCP_PROJECT) + table_ref = client.dataset(_BIGQUERY_DATASET).table(_BIGQUERY_TABLE) + table_obj = client.get_table(table_ref) + errors = client.insert_rows(table_obj, [row]) + if errors: + raise RuntimeError( + f'Failed to insert row into MLCompass table: {str(errors)}') + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Export datapoint to MLCompass") + parser.add_argument( + "--result_file", + type=str, + help="Path to result file generated by report_result.sh", + ) + args = parser.parse_args() + metrics = read_metrics_file(args.result_file) + export(metrics) + + +if __name__ == "__main__": + main() diff --git a/.buildkite/benchmark/scripts/report_result.sh b/.buildkite/benchmark/scripts/report_result.sh index 208f3bac94..130cb3ff2d 100755 --- a/.buildkite/benchmark/scripts/report_result.sh +++ b/.buildkite/benchmark/scripts/report_result.sh @@ -375,3 +375,10 @@ else echo "Warning: $RESULT_FILE not found. No results to display." fi fi + +if [[ "${MLCOMPASS_EXPORT_ENABLED:-false}" == "true" ]]; then + echo "--- Reporting to MLCompass" + python3 "$(dirname "${BASH_SOURCE[0]}")/mlcompass_export.py" --result_file="$RESULT_FILE" +else + echo "--- Reporting to MLCompass (skipped)" +fi diff --git a/.buildkite/benchmark/scripts/run_job.sh b/.buildkite/benchmark/scripts/run_job.sh index c0557bda02..bda24d3fbd 100755 --- a/.buildkite/benchmark/scripts/run_job.sh +++ b/.buildkite/benchmark/scripts/run_job.sh @@ -102,7 +102,14 @@ declare -a BENCHMARK_DOCKER_ARGS=( "-e" "BUILDKITE=${BUILDKITE}" "-e" "BUILDKITE_AGENT_NAME=${BUILDKITE_AGENT_NAME}" "-e" "BUILDKITE_AGENT_META_DATA_QUEUE=${BUILDKITE_AGENT_META_DATA_QUEUE}" + "-e" "BUILDKITE_BUILD_NUMBER=${BUILDKITE_BUILD_NUMBER}" + "-e" "BUILDKITE_JOB_ID=${BUILDKITE_JOB_ID}" "-e" "UPLOAD_DB=${UPLOAD_DB:-true}" + "-e" "MLCOMPASS_EXECUTION_MODE=${MLCOMPASS_EXECUTION_MODE:-}" + "-e" "MLCOMPASS_EXPORT_ENABLED=${MLCOMPASS_EXPORT_ENABLED:-}" + "-e" "MLCOMPASS_TEST_NAME=${MLCOMPASS_TEST_NAME:-}" + "-e" "MLCOMPASS_TRACKING_ID=${MLCOMPASS_TRACKING_ID:-}" + "-e" "MLCOMPASS_SPONGE_ID=${MLCOMPASS_SPONGE_ID:-}" ) BENCHMARK_DOCKER_ARGS_STR="$(printf '%s\n' "${BENCHMARK_DOCKER_ARGS[@]}")" diff --git a/requirements_benchmarking.txt b/requirements_benchmarking.txt index 2d46f13324..3d67d78c3c 100644 --- a/requirements_benchmarking.txt +++ b/requirements_benchmarking.txt @@ -7,3 +7,4 @@ rouge-score scikit-learn pandas gdown<6.0.0 +google-cloud-bigquery From d37f0fee2678d1d01c9026c7470404f11b4d1c80 Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Tue, 9 Jun 2026 09:20:01 +0000 Subject: [PATCH 32/37] [skip ci] Update nightly support matrices (v6e/v7x) Signed-off-by: Buildkite Bot --- support_matrices/nightly/v7x/default/model_support_matrix.csv | 2 +- .../nightly/v7x/default/parallelism_support_matrix.csv | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/support_matrices/nightly/v7x/default/model_support_matrix.csv b/support_matrices/nightly/v7x/default/model_support_matrix.csv index 11969b37a2..20a7a25233 100644 --- a/support_matrices/nightly/v7x/default/model_support_matrix.csv +++ b/support_matrices/nightly/v7x/default/model_support_matrix.csv @@ -23,7 +23,7 @@ Model,Type,UnitTest,Accuracy/Correctness,Benchmark "google/gemma-3-27b-it",Text,✅ Passing,✅ Passing,✅ Passing "meta-llama/Llama-3.1-8B-Instruct",Text,✅ Passing,✅ Passing,✅ Passing "meta-llama/Llama-3.3-70B-Instruct",Text,✅ Passing,✅ Passing,✅ Passing -"moonshotai/Kimi-K2-Thinking",Text,❓ Untested,❓ Untested,❓ Untested +"moonshotai/Kimi-K2-Thinking",Text,❌ Failing,❓ Untested,❓ Untested "moonshotai/Kimi-K2.6",Text,✅ Passing,❓ Untested,❓ Untested "openai/gpt-oss-120b",Text,✅ Passing,✅ Passing,❓ Untested "openai/gpt-oss-20b",Text,❓ Untested,❓ Untested,❓ Untested diff --git a/support_matrices/nightly/v7x/default/parallelism_support_matrix.csv b/support_matrices/nightly/v7x/default/parallelism_support_matrix.csv index a6cc9507ed..5e1ccf6dc6 100644 --- a/support_matrices/nightly/v7x/default/parallelism_support_matrix.csv +++ b/support_matrices/nightly/v7x/default/parallelism_support_matrix.csv @@ -1,7 +1,7 @@ Feature,Single-Host CorrectnessTest,Single-Host PerformanceTest,Multi-Host CorrectnessTest,Multi-Host PerformanceTest "CP",❓ Untested,❓ Untested,❓ Untested,❓ Untested "DP",✅ Passing,✅ Passing,❓ Untested,❓ Untested -"EP",✅ Passing,✅ Passing,❓ Untested,❓ Untested +"EP",❌ Failing,❓ Untested,❓ Untested,❓ Untested "PP",✅ Passing,✅ Passing,✅ Passing,✅ Passing "SP",✅ Passing,❓ Untested,❓ Untested,❓ Untested -"TP",✅ Passing,✅ Passing,❓ Untested,❓ Untested +"TP",✅ Passing,❌ Failing,❓ Untested,❓ Untested From 444175fe0262c6e05695001192bb41bca7391ba6 Mon Sep 17 00:00:00 2001 From: Teresa Chen <120631815+boe20211@users.noreply.github.com> Date: Tue, 9 Jun 2026 17:30:40 +0800 Subject: [PATCH 33/37] Add Deepseek benchmark case (#2667) Signed-off-by: Teresa Chen --- .../benchmark/cases/daily/deepseek.json | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 .buildkite/benchmark/cases/daily/deepseek.json diff --git a/.buildkite/benchmark/cases/daily/deepseek.json b/.buildkite/benchmark/cases/daily/deepseek.json new file mode 100644 index 0000000000..466e9af3a9 --- /dev/null +++ b/.buildkite/benchmark/cases/daily/deepseek.json @@ -0,0 +1,67 @@ +{ + "global_env": { + "GCP_PROJECT_ID": "cloud-tpu-inference-test", + "GCS_BUCKET": "vllm-cb-storage2", + "GCP_INSTANCE_ID": "vllm-bm-inst", + "GCP_DATABASE_ID": "vllm-bm-bk-runs" + }, + "benchmark_cases": [ + { + "case_name": "deepseek-r1-dataset_custom-inlen_1024-outlen_4096-prelen_0", + "ci_queue": [ + "tpu_v7x_8_queue" + ], + "env": { + "RUN_TYPE": "DAILY", + "MODELTAG": "NEW", + "EXPECTED_ETEL": 43200000, + "INPUT_LEN": 1024, + "OUTPUT_LEN": 4096, + "PREFIX_LEN": 0 + }, + "server_command_options": { + "command_type": "vllm_serve", + "env": { + "VLLM_USE_V1": "1", + "MODEL_IMPL_TYPE": "vllm", + "NEW_MODEL_DESIGN": "1", + "TPU_BACKEND_TYPE": "jax", + "VLLM_MLA_DISABLE": "0", + "MOE_REQUANTIZE_BLOCK_SIZE": "512", + "MOE_REQUANTIZE_WEIGHT_DTYPE": "fp4" + }, + "args": { + "model": "deepseek-ai/DeepSeek-R1", + "seed": 42, + "max-num-seqs": 16, + "max-num-batched-tokens": 1024, + "tensor-parallel-size": { + "v7x-8": 8 + }, + "max-model-len": 5120, + "kv-cache-dtype": "fp8", + "gpu-memory-utilization": 0.95, + "enable-expert-parallel": true, + "no-enable-prefix-caching": true, + "async-scheduling": true, + "additional-config": "{\"sharding\": {\"sharding_strategy\": {\"enable_dp_attention\": true}}}" + } + }, + "client_command_options": { + "command_type": "vllm_bench_serve", + "args": { + "model": "deepseek-ai/DeepSeek-R1", + "backend": "vllm", + "request-rate": "inf", + "dataset-name": "custom", + "dataset-path": "/workspace/tpu_inference/artifacts/dataset/DeepSeek-R1/inlen1024_outlen4096_prefixlen0.jsonl", + "num-prompts": 64, + "percentile-metrics": "ttft,tpot,itl,e2el", + "ignore-eos": true, + "custom-output-len": 4096, + "skip-chat-template": true + } + } + } + ] +} \ No newline at end of file From e87ca8d78e2ddb4e2f8d10689f89ef3e0ad35844 Mon Sep 17 00:00:00 2001 From: Buildkite Bot Date: Tue, 9 Jun 2026 10:57:12 +0000 Subject: [PATCH 34/37] [skip ci] Update nightly support matrices for vllm (v6e/v7x) Signed-off-by: Buildkite Bot --- support_matrices/nightly/v7x/vllm/feature_support_matrix.csv | 2 +- support_matrices/nightly/v7x/vllm/model_support_matrix.csv | 2 +- .../nightly/v7x/vllm/parallelism_support_matrix.csv | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv b/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv index 21b2ea4ef7..3cada0d1fe 100644 --- a/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv +++ b/support_matrices/nightly/v7x/vllm/feature_support_matrix.csv @@ -13,6 +13,6 @@ Feature,CorrectnessTest,PerformanceTest "async scheduler",✅ Passing,✅ Passing "hybrid kv cache",✅ Passing,❓ Untested "multi-host",✅ Passing,❓ Untested -"runai_model_streamer_loader",✅ Passing,❓ Untested +"runai_model_streamer_loader",❌ Failing,❓ Untested "sampling_params",✅ Passing,❓ Untested "structured_decoding",✅ Passing,❓ Untested diff --git a/support_matrices/nightly/v7x/vllm/model_support_matrix.csv b/support_matrices/nightly/v7x/vllm/model_support_matrix.csv index 4d3b3a2104..6e58ab7a48 100644 --- a/support_matrices/nightly/v7x/vllm/model_support_matrix.csv +++ b/support_matrices/nightly/v7x/vllm/model_support_matrix.csv @@ -17,7 +17,7 @@ Model,Type,UnitTest,Accuracy/Correctness,Benchmark "Qwen/Qwen3.5-397B-A17B",Text,✅ Passing,✅ Passing,✅ Passing "deepseek-ai/DeepSeek-Math-V2",Text,❓ Untested,❓ Untested,❓ Untested "deepseek-ai/DeepSeek-R1",Text,✅ Passing,✅ Passing,❓ Untested -"deepseek-ai/DeepSeek-V3.1",Text,❓ Untested,❓ Untested,❓ Untested +"deepseek-ai/DeepSeek-V3.1",Text,❓ Untested,❌ Failing,❓ Untested "deepseek-ai/DeepSeek-V3.2",Text,❓ Untested,❓ Untested,❓ Untested "deepseek-ai/DeepSeek-V3.2-Speciale",Text,❓ Untested,❓ Untested,❓ Untested "google/gemma-3-27b-it",Text,✅ Passing,✅ Passing,✅ Passing diff --git a/support_matrices/nightly/v7x/vllm/parallelism_support_matrix.csv b/support_matrices/nightly/v7x/vllm/parallelism_support_matrix.csv index a6cc9507ed..f5c1547228 100644 --- a/support_matrices/nightly/v7x/vllm/parallelism_support_matrix.csv +++ b/support_matrices/nightly/v7x/vllm/parallelism_support_matrix.csv @@ -4,4 +4,4 @@ Feature,Single-Host CorrectnessTest,Single-Host PerformanceTest,Multi-Host Corre "EP",✅ Passing,✅ Passing,❓ Untested,❓ Untested "PP",✅ Passing,✅ Passing,✅ Passing,✅ Passing "SP",✅ Passing,❓ Untested,❓ Untested,❓ Untested -"TP",✅ Passing,✅ Passing,❓ Untested,❓ Untested +"TP",✅ Passing,❌ Failing,❓ Untested,❓ Untested From 862918b2ad16528bf49839f91379a759bd57a206 Mon Sep 17 00:00:00 2001 From: Jacob Platin <31421084+jrplatin@users.noreply.github.com> Date: Tue, 9 Jun 2026 10:46:44 -0700 Subject: [PATCH 35/37] [Kimi-K2.6] Add correctness tests for LM only (#2848) Signed-off-by: Signed-off-by: Jacob Platin --- .buildkite/models/moonshotai_Kimi-K2_6.yml | 23 +++++++++++++++++----- tests/e2e/benchmarking/mmlu.sh | 13 +++++++++--- 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/.buildkite/models/moonshotai_Kimi-K2_6.yml b/.buildkite/models/moonshotai_Kimi-K2_6.yml index b62d1f5ae9..833432a828 100644 --- a/.buildkite/models/moonshotai_Kimi-K2_6.yml +++ b/.buildkite/models/moonshotai_Kimi-K2_6.yml @@ -54,14 +54,27 @@ steps: depends_on: "${TPU_VERSION:-tpu6e}_record_moonshotai_Kimi-K2_6_UnitTest" soft_fail: true agents: - queue: cpu # TODO: replace with execution environment + queue: ${TPU_QUEUE_MULTI:-tpu_v6e_8_queue} env: - TEST_MODEL: moonshotai/Kimi-K2.6 - TENSOR_PARALLEL_SIZE: "${TENSOR_PARALLEL_SIZE_SINGLE:-1}" - MINIMUM_ACCURACY_THRESHOLD: 0 # TODO : replace 0 with your accuracy threshold + NEW_MODEL_DESIGN: "1" + USE_V6E8_QUEUE: "False" + TPU_VERSION: "tpu7x" + SKIP_ACCURACY_TESTS: "False" + VLLM_MLA_DISABLE: "0" + MOE_REQUANTIZE_BLOCK_SIZE: "512" + MOE_REQUANTIZE_WEIGHT_DTYPE: "fp4" + DEVICE_COUNT: "8" + MODEL_IMPL_TYPE: "vllm" + USE_V7X8_QUEUE: "True" + MINIMUM_ACCURACY_THRESHOLD: "0.86" commands: - | - buildkite-agent meta-data set "${TPU_VERSION:-tpu6e}_moonshotai_Kimi-K2_6_Accuracy" "unverified" + if [ "${TPU_VERSION:-tpu6e}" = "tpu6e" ]; then + buildkite-agent meta-data set "${TPU_VERSION:-tpu6e}_moonshotai_Kimi-K2_6_Accuracy" "not enough HBM" + else + .buildkite/scripts/run_in_docker.sh \ + bash /workspace/tpu_inference/tests/e2e/benchmarking/mmlu.sh -m "gs://tpu-commons-ci/moonshootai/kimi/2.6" -n 14042 -l 4 + fi - label: "${TPU_VERSION:-tpu6e} Record accuracy result for moonshotai/Kimi-K2_6" key: "${TPU_VERSION:-tpu6e}_record_moonshotai_Kimi-K2_6_Accuracy" depends_on: "${TPU_VERSION:-tpu6e}_moonshotai_Kimi-K2_6_Accuracy" diff --git a/tests/e2e/benchmarking/mmlu.sh b/tests/e2e/benchmarking/mmlu.sh index 5a7dac48bc..3d97f9cee0 100644 --- a/tests/e2e/benchmarking/mmlu.sh +++ b/tests/e2e/benchmarking/mmlu.sh @@ -263,8 +263,8 @@ for model_name in $model_list; do elif [ "$TPU_VERSION" == "tpu7x" ]; then # Set the default value to 2 for tpu v7x for most models current_device_count=2 - # If using large model (e.g. DeepSeek, then set the count to 8) - if [[ "${model_name,,}" == *"deepseek"* ]]; then + # If using large model (e.g. DeepSeek or Kimi, then set the count to 8) + if [[ "${model_name,,}" == *"deepseek"* ]] || [[ "${model_name,,}" == *"kimi"* ]]; then current_device_count=8 fi else @@ -295,10 +295,17 @@ for model_name in $model_list; do current_serve_args+=(--enable-expert-parallel) current_serve_args+=(--additional_config '{"sharding": {"sharding_strategy": {"enable_dp_attention": true}}}') elif [ "$MODEL_IMPL_TYPE" == "flax_nnx" ]; then - current_serve_args+=(--additional_config '{"sharding": {"sharding_strategy": + current_serve_args+=(--additional_config '{"sharding": {"sharding_strategy": {"enable_dp_attention": true, "expert_parallelism": '"${current_device_count}"', "tensor_parallelism": 1}}, "replicate_attn_weights": "True", "sparse_matmul": "True"}') fi + elif [[ "${model_name,,}" == *"kimi"* ]]; then + export TIMEOUT_SECONDS=3600 # Kimi needs a longer timeout + max_batched_tokens=1024 + served_name=moonshotai/Kimi-K2.6 + TARGET_ACCURACY="0.87" + current_serve_args+=(--kv-cache-dtype=fp8 --gpu-memory-utilization 0.977 --limit-mm-per-prompt='{"image": 0, "video": 0, "vision_chunk": 0}' ) + current_serve_args+=(--served-model-name "${served_name}" --load-format=runai_streamer --enable-expert-parallel --trust-remote-code --kv-cache-dtype=fp8 --additional-config '{"sharding": {"sharding_strategy": {"enable_dp_attention": true, "tensor_parallelism": '"${current_device_count}"'}}}') fi fi From d595ee4118d2c82b42ab4e4ef4a0d3492a9c0971 Mon Sep 17 00:00:00 2001 From: Charles Li Date: Mon, 8 Jun 2026 23:03:28 +0000 Subject: [PATCH 36/37] Add Combo test cases --- .../combo/test_combo_chunked_offload_spec.py | 202 ++++++++++++++++++ .../test_combo_hybrid_kv_cache_prefix.py | 117 ++++++++++ .../combo/test_combo_offload_prefix_spec.py | 181 ++++++++++++++++ .../test_combo_sampling_params_prefix.py | 107 ++++++++++ 4 files changed, 607 insertions(+) create mode 100644 tests/e2e/combo/test_combo_chunked_offload_spec.py create mode 100644 tests/e2e/combo/test_combo_hybrid_kv_cache_prefix.py create mode 100644 tests/e2e/combo/test_combo_offload_prefix_spec.py create mode 100644 tests/e2e/combo/test_combo_sampling_params_prefix.py diff --git a/tests/e2e/combo/test_combo_chunked_offload_spec.py b/tests/e2e/combo/test_combo_chunked_offload_spec.py new file mode 100644 index 0000000000..6df9f2da82 --- /dev/null +++ b/tests/e2e/combo/test_combo_chunked_offload_spec.py @@ -0,0 +1,202 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import itertools +import os +import time +from typing import Optional, Union + +import pytest +from vllm import LLM, SamplingParams +from vllm.config import KVTransferConfig + + +def parse_outputs(outputs): + output_token_ids = [] + generated_texts = [] + for output in outputs: + completion = output.outputs[0] + generated_text = completion.text + token_ids = completion.token_ids + generated_texts.append(generated_text) + output_token_ids.append(token_ids) + return generated_texts, output_token_ids + + +def get_sampling_config(): + """deterministic sampling config""" + return SamplingParams(temperature=0.0, + max_tokens=20, + seed=42, + ignore_eos=True) + + +def get_kv_transfer_config(): + """use TPUOffloadConnector""" + return KVTransferConfig( + kv_connector="TPUOffloadConnector", + kv_role="kv_both", + kv_connector_module_path="tpu_inference.offload.tpu_offload_connector", + ) + + +def _test_combo_chunked_offload_spec_accuracy( + monkeypatch: pytest.MonkeyPatch, + model_name: str, + speculative_config: Optional[Union[dict, str]], + enable_chunked_prefill: bool, + async_scheduling: Union[bool, str] = "auto", + cpu_chunks: str = "8", + max_output_len: Optional[int] = None, +): + sampling_config = get_sampling_config() + if max_output_len is not None: + sampling_config.max_tokens = max_output_len + kv_transfer_config = get_kv_transfer_config() + + prompts = [ + "The quick brown fox jumps over the lazy dog.", + "Artificial intelligence is transformational for scientific research and computing.", + ] + num_requests = len(prompts) + + with monkeypatch.context(): + # Standard environmental configuration for JAX-TPU KV Offloading + os.environ['SKIP_JAX_PRECOMPILE'] = '0' + os.environ['TPU_OFFLOAD_SKIP_JAX_PRECOMPILE'] = '0' + os.environ['TPU_OFFLOAD_DECODE_SAVE'] = '1' + os.environ['TPU_OFFLOAD_BATCHED_SAVE'] = '0' + os.environ['TPU_OFFLOAD_NUM_CPU_CHUNKS'] = cpu_chunks + + # Ensure Hugging Face respects offline cache if specified + os.environ["HF_HOME"] = os.environ.get("HF_HOME", "~/.cache/huggingface") + os.environ["HF_HUB_CACHE"] = os.environ.get("HF_HUB_CACHE", "~/.cache/huggingface/hub") + + tensor_parallel_size = int(os.environ.get("TPU_TP_SIZE", "8")) + + if async_scheduling == "auto": + async_scheduling_val = not speculative_config + else: + async_scheduling_val = bool(async_scheduling) + + llm = LLM( + model=model_name, + max_model_len=1024, + async_scheduling=async_scheduling_val, + tensor_parallel_size=tensor_parallel_size, + enable_chunked_prefill=enable_chunked_prefill, # Chunked Prefill Switch + kv_transfer_config=kv_transfer_config, # F03: KV Cache Offloading + speculative_config=speculative_config, # F10: Speculative Decoding + ) + + # --- Pass 1: Cold Generation (Calculates and offloads KV Cache) --- + print(f"\n--- Pass 1: Generating for {num_requests} requests ---") + t0 = time.time() + outputs1 = llm.generate(prompts, sampling_config) + pass1_time = time.time() - t0 + print(f"Pass 1 generation completed in {pass1_time:.4f} seconds") + out_texts1, out_tokens1 = parse_outputs(outputs1) + time.sleep(5) + + # --- Resetting prefix cache in TPU HBM --- + # Forces next pass to load prefix KV cache from Host CPU DRAM instead of recalculating + print("\n--- Resetting prefix cache (evicting from TPU HBM) ---") + llm.llm_engine.engine_core.reset_prefix_cache() + time.sleep(2) + + # --- Pass 2: Warm Generation (Loads KV Cache from CPU DRAM) --- + print(f"\n--- Pass 2: Generating again (should load from CPU DRAM) ---") + t0 = time.time() + outputs2 = llm.generate(prompts, sampling_config) + pass2_time = time.time() - t0 + print(f"Pass 2 generation completed in {pass2_time:.4f} seconds") + out_texts2, out_tokens2 = parse_outputs(outputs2) + time.sleep(1) + + print("\n" + "=" * 80) + print("Accuracy Comparison Results") + print(f"Pass 1 generate time: {pass1_time * 1000:.2f} ms") + print(f"Pass 2 generate time: {pass2_time * 1000:.2f} ms") + print("=" * 80) + for i in range(len(out_texts1)): + print(f"\nRequest {i}:") + print(f" Pass 1 Output Text: {out_texts1[i]!r}") + print(f" Pass 2 Output Text: {out_texts2[i]!r}") + print(f" Pass 1 Output Tokens: {out_tokens1[i]}") + print(f" Pass 2 Output Tokens: {out_tokens2[i]}") + print("\n" + "=" * 80) + + # Output 1 and Output 2 must be bit-for-bit identical + assert len(out_texts1) == len(out_texts2) + assert len(out_tokens1) == len(out_tokens2) + for i in range(len(out_texts1)): + assert out_texts1[i] == out_texts2[i], f"Text mismatch in request {i}" + assert out_tokens1[i] == out_tokens2[i], f"Token mismatch in request {i}" + + del llm + # Waiting for TPUs to be released. + time.sleep(10) + + +@pytest.mark.parametrize("enable_chunked_prefill", [True, False]) +@pytest.mark.parametrize("async_scheduling", [ + pytest.param(True, marks=pytest.mark.xfail(reason="Async scheduling is not supported with speculative decoding on TPU backend yet", strict=False)), + pytest.param(False) +]) +def test_combo_chunked_ngram_llama_3b(monkeypatch: pytest.MonkeyPatch, enable_chunked_prefill: bool, async_scheduling: bool): + """Tests Llama-3.2-3B-Instruct using Ngram Speculative Decoding, Chunked Prefill, and KV Offloading""" + model_name = os.environ.get("MODEL_NAME", "meta-llama/Llama-3.2-3B-Instruct") + speculative_config = { + "method": "ngram", + "prompt_lookup_max": 2, + "prompt_lookup_min": 2, + "num_speculative_tokens": 4, + } + _test_combo_chunked_offload_spec_accuracy( + monkeypatch=monkeypatch, + model_name=model_name, + speculative_config=speculative_config, + enable_chunked_prefill=enable_chunked_prefill, + async_scheduling=async_scheduling, + cpu_chunks="8", + ) + + +@pytest.mark.parametrize("enable_chunked_prefill", [True, False]) +@pytest.mark.parametrize("async_scheduling", [ + pytest.param(True, marks=pytest.mark.xfail(reason="Async scheduling is not supported with speculative decoding on TPU backend yet", strict=False)), + pytest.param(False) +]) +def test_combo_chunked_eagle3_llama_8b(monkeypatch: pytest.MonkeyPatch, enable_chunked_prefill: bool, async_scheduling: bool): + """Tests Llama-3 8B using Eagle3 Speculative Decoding, Chunked Prefill, and KV Offloading""" + model_name = os.environ.get("MODEL_NAME", "meta-llama/Meta-Llama-3.1-8B-Instruct") + draft_model = os.environ.get("DRAFT_MODEL_NAME", "unkmaster/EAGLE3-LLaMA3.1-Instruct-8B") + + model_impl = os.environ.get("MODEL_IMPL_TYPE", "auto") + monkeypatch.setenv("DRAFT_MODEL_IMPL_TYPE", model_impl) + + speculative_config = { + "method": "eagle3", + "model": draft_model, + "num_speculative_tokens": 3, + "draft_tensor_parallel_size": 1, + } + _test_combo_chunked_offload_spec_accuracy( + monkeypatch=monkeypatch, + model_name=model_name, + speculative_config=speculative_config, + enable_chunked_prefill=enable_chunked_prefill, + async_scheduling=async_scheduling, + cpu_chunks="8", + ) diff --git a/tests/e2e/combo/test_combo_hybrid_kv_cache_prefix.py b/tests/e2e/combo/test_combo_hybrid_kv_cache_prefix.py new file mode 100644 index 0000000000..152c41eb54 --- /dev/null +++ b/tests/e2e/combo/test_combo_hybrid_kv_cache_prefix.py @@ -0,0 +1,117 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import time +from typing import List, Tuple + +import pytest +from vllm import LLM, SamplingParams + + +MODEL_NAME = "google/gemma-3-4b-it" + + +def _parse_outputs(outputs) -> Tuple[List[str], List[Tuple[int, ...]]]: + texts = [] + token_ids = [] + for output in outputs: + completion = output.outputs[0] + texts.append(completion.text) + token_ids.append(tuple(completion.token_ids)) + return texts, token_ids + + +def _reset_engine_prefix_cache(llm: LLM) -> None: + if hasattr(llm, "reset_prefix_cache"): + llm.reset_prefix_cache() + return + llm.llm_engine.engine_core.reset_prefix_cache() + + +@pytest.fixture +def sampling_params(): + return SamplingParams( + temperature=0.0, + max_tokens=16, + seed=42, + ignore_eos=True, + ) + + +@pytest.fixture +def shared_prefix_prompts(): + shared_prefix = ( + "A family is planning a weekend trip to a quiet town near the coast. " + "They want to visit a local museum, walk through the old market, try " + "a seafood restaurant, and spend some time watching the sunset by the " + "water. ") + return [ + shared_prefix + "Write a short travel plan for Saturday.", + shared_prefix + "Suggest what they should pack for the trip.", + ] + + +def test_kv_cache_prefix_caching_with_hybrid_kv_cache( + monkeypatch: pytest.MonkeyPatch, + sampling_params: SamplingParams, + shared_prefix_prompts: List[str], +): + """ + Exercise TPU Prefix Caching combined with Hybrid KV Cache (HMA). + """ + monkeypatch.setenv("MODEL_IMPL_TYPE", "vllm") + monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "0") + + llm = LLM( + model=MODEL_NAME, + max_model_len=192, + tensor_parallel_size=8, + max_num_batched_tokens=2048, + max_num_seqs=64, + enable_prefix_caching=True, + disable_hybrid_kv_cache_manager=False, # Enable Hybrid KV Cache Manager + ) + + try: + # Step 1: Warm up and populate the in-memory prefix cache + outputs1 = llm.generate(shared_prefix_prompts, sampling_params) + texts1, tokens1 = _parse_outputs(outputs1) + + # Step 2: Repeat generations to verify prefix cache hits + outputs2 = llm.generate(shared_prefix_prompts, sampling_params) + texts2, tokens2 = _parse_outputs(outputs2) + + assert texts1 == texts2 + assert tokens1 == tokens2 + + # Step 3: Clear the in-memory prefix cache index + _reset_engine_prefix_cache(llm) + time.sleep(1) + + filler_prompts = [ + "Explain quantum computing in simple terms.", + "Write a short note about deterministic sampling.", + ] + llm.generate(filler_prompts, sampling_params) + + # Step 4: Verify outputs match identically after recomputation + outputs3 = llm.generate(shared_prefix_prompts, sampling_params) + texts3, tokens3 = _parse_outputs(outputs3) + + assert texts1 == texts3 + assert tokens1 == tokens3 + finally: + if 'llm' in locals() and hasattr(llm.llm_engine, "shutdown"): + llm.llm_engine.shutdown() + time.sleep(5) \ No newline at end of file diff --git a/tests/e2e/combo/test_combo_offload_prefix_spec.py b/tests/e2e/combo/test_combo_offload_prefix_spec.py new file mode 100644 index 0000000000..6d04e7c94e --- /dev/null +++ b/tests/e2e/combo/test_combo_offload_prefix_spec.py @@ -0,0 +1,181 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import itertools +import os +import time +from typing import Optional, Union + +import pytest +from vllm import LLM, SamplingParams +from vllm.config import KVTransferConfig + + +def parse_outputs(outputs): + output_token_ids = [] + generated_texts = [] + for output in outputs: + completion = output.outputs[0] + generated_text = completion.text + token_ids = completion.token_ids + generated_texts.append(generated_text) + output_token_ids.append(token_ids) + return generated_texts, output_token_ids + + +def get_sampling_config(): + """deterministic sampling config""" + return SamplingParams(temperature=0.0, + max_tokens=20, + seed=42, + ignore_eos=True) + + +def get_kv_transfer_config(): + """use TPUOffloadConnector""" + return KVTransferConfig( + kv_connector="TPUOffloadConnector", + kv_role="kv_both", + kv_connector_module_path="tpu_inference.offload.tpu_offload_connector", + ) + + +def _test_combo_offload_prefix_spec_accuracy( + monkeypatch: pytest.MonkeyPatch, + model_name: str, + speculative_config: Optional[Union[dict, str]], + cpu_chunks: str = "8", + max_output_len: Optional[int] = None, +): + sampling_config = get_sampling_config() + if max_output_len is not None: + sampling_config.max_tokens = max_output_len + kv_transfer_config = get_kv_transfer_config() + + prompts = [ + "The quick brown fox jumps over the lazy dog.", + "Artificial intelligence is transformational for scientific research and computing.", + ] + num_requests = len(prompts) + + with monkeypatch.context(): + # Standard environmental configuration for JAX-TPU KV Offloading + os.environ['SKIP_JAX_PRECOMPILE'] = '0' + os.environ['TPU_OFFLOAD_SKIP_JAX_PRECOMPILE'] = '0' + os.environ['TPU_OFFLOAD_DECODE_SAVE'] = '1' + os.environ['TPU_OFFLOAD_BATCHED_SAVE'] = '0' + os.environ['TPU_OFFLOAD_NUM_CPU_CHUNKS'] = cpu_chunks + + # Ensure Hugging Face respects offline cache if specified + os.environ["HF_HOME"] = os.environ.get("HF_HOME", "~/.cache/huggingface") + os.environ["HF_HUB_CACHE"] = os.environ.get("HF_HUB_CACHE", "~/.cache/huggingface/hub") + + tensor_parallel_size = int(os.environ.get("TPU_TP_SIZE", "8")) + + llm = LLM( + model=model_name, + max_model_len=1024, + async_scheduling=not speculative_config, + tensor_parallel_size=tensor_parallel_size, + enable_prefix_caching=True, + kv_transfer_config=kv_transfer_config, + speculative_config=speculative_config, + ) + + # --- Pass 1: Cold Generation (Calculates and offloads KV Cache) --- + print(f"\n--- Pass 1: Generating for {num_requests} requests ---") + t0 = time.time() + outputs1 = llm.generate(prompts, sampling_config) + pass1_time = time.time() - t0 + print(f"Pass 1 generation completed in {pass1_time:.4f} seconds") + out_texts1, out_tokens1 = parse_outputs(outputs1) + time.sleep(5) + + # --- Resetting prefix cache in TPU HBM --- + # Forces next pass to load prefix KV cache from Host CPU DRAM instead of recalculating + print("\n--- Resetting prefix cache (evicting from TPU HBM) ---") + llm.llm_engine.engine_core.reset_prefix_cache() + time.sleep(2) + + # --- Pass 2: Warm Generation (Loads KV Cache from CPU DRAM) --- + print(f"\n--- Pass 2: Generating again (should load from CPU DRAM) ---") + t0 = time.time() + outputs2 = llm.generate(prompts, sampling_config) + pass2_time = time.time() - t0 + print(f"Pass 2 generation completed in {pass2_time:.4f} seconds") + out_texts2, out_tokens2 = parse_outputs(outputs2) + time.sleep(1) + + print("\n" + "=" * 80) + print("Accuracy Comparison Results") + print(f"Pass 1 generate time: {pass1_time * 1000:.2f} ms") + print(f"Pass 2 generate time: {pass2_time * 1000:.2f} ms") + print("=" * 80) + for i in range(len(out_texts1)): + print(f"\nRequest {i}:") + print(f" Pass 1 Output Text: {out_texts1[i]!r}") + print(f" Pass 2 Output Text: {out_texts2[i]!r}") + print(f" Pass 1 Output Tokens: {out_tokens1[i]}") + print(f" Pass 2 Output Tokens: {out_tokens2[i]}") + print("\n" + "=" * 80) + + # Output 1 and Output 2 must be bit-for-bit identical + assert len(out_texts1) == len(out_texts2) + assert len(out_tokens1) == len(out_tokens2) + for i in range(len(out_texts1)): + assert out_texts1[i] == out_texts2[i], f"Text mismatch in request {i}" + assert out_tokens1[i] == out_tokens2[i], f"Token mismatch in request {i}" + + del llm + # Waiting for TPUs to be released. + time.sleep(10) + + +def test_combo_ngram_llama_3b(monkeypatch: pytest.MonkeyPatch): + """Tests Llama-3.2-3B-Instruct using Ngram Speculative Decoding, Prefix Caching, and KV Offloading""" + model_name = os.environ.get("MODEL_NAME", "meta-llama/Llama-3.2-3B-Instruct") + speculative_config = { + "method": "ngram", + "prompt_lookup_max": 2, + "prompt_lookup_min": 2, + "num_speculative_tokens": 4, + } + _test_combo_offload_prefix_spec_accuracy( + monkeypatch=monkeypatch, + model_name=model_name, + speculative_config=speculative_config, + cpu_chunks="8", + ) + + +def test_combo_eagle3_llama_8b(monkeypatch: pytest.MonkeyPatch): + """Tests Llama-3 8B using Eagle3 Speculative Decoding, Prefix Caching, and KV Offloading""" + model_name = os.environ.get("MODEL_NAME", "meta-llama/Meta-Llama-3.1-8B-Instruct") + draft_model = os.environ.get("DRAFT_MODEL_NAME", "unkmaster/EAGLE3-LLaMA3.1-Instruct-8B") + + model_impl = os.environ.get("MODEL_IMPL_TYPE", "auto") + monkeypatch.setenv("DRAFT_MODEL_IMPL_TYPE", model_impl) + + speculative_config = { + "method": "eagle3", + "model": draft_model, + "num_speculative_tokens": 3, + "draft_tensor_parallel_size": 1, + } + _test_combo_offload_prefix_spec_accuracy( + monkeypatch=monkeypatch, + model_name=model_name, + speculative_config=speculative_config, + cpu_chunks="8", + ) diff --git a/tests/e2e/combo/test_combo_sampling_params_prefix.py b/tests/e2e/combo/test_combo_sampling_params_prefix.py new file mode 100644 index 0000000000..7d2eb5c603 --- /dev/null +++ b/tests/e2e/combo/test_combo_sampling_params_prefix.py @@ -0,0 +1,107 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from vllm import LLM, SamplingParams + +# Global constant replacing class attribute +SYSTEM_CONTEXT = ( + "You are a helpful and precise assistant. " + "The quick brown fox jumps over the lazy dog. " + "Artificial Intelligence and TPU acceleration make inference incredibly fast." +) + + +@pytest.fixture(scope="module") +def tpu_llm(): + """Initialize TPU vLLM instance with Prefix Caching enabled.""" + return LLM( + model='meta-llama/Llama-3.2-1B-Instruct', + max_model_len=1024, + max_num_seqs=4, + enable_prefix_caching=True, + ) + + +def _apply_template(llm: LLM, system_msg: str, user_msg: str) -> str: + """Helper to apply proper Chat Template for Llama-3-Instruct models.""" + tokenizer = llm.get_tokenizer() + messages = [ + {"role": "system", "content": system_msg}, + {"role": "user", "content": user_msg} + ] + return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + +def test_cache_hit_with_deterministic_sampling(tpu_llm: LLM): + """Verify that greedy sampling (temp=0) remains deterministic after a cache hit.""" + prompt1 = _apply_template(tpu_llm, SYSTEM_CONTEXT, "What is 5 + 5? Answer with just the number:") + prompt2 = _apply_template(tpu_llm, SYSTEM_CONTEXT, "What is 5 + 5? Answer with just the number:") + + sampling_params = SamplingParams(temperature=0, max_tokens=5) + + outputs1 = tpu_llm.generate([prompt1], sampling_params) + outputs2 = tpu_llm.generate([prompt2], sampling_params) + + text1 = outputs1[0].outputs[0].text.strip() + text2 = outputs2[0].outputs[0].text.strip() + + assert text1 == text2, f"Cache hit altered greedy output: '{text1}' vs '{text2}'" + + +def test_cache_hit_with_random_sampling(tpu_llm: LLM): + """Verify that random sampling (temp > 0) maintains diversity on cache hits.""" + prompt = _apply_template(tpu_llm, SYSTEM_CONTEXT, "Write a completely random and creative single word:") + sampling_params = SamplingParams(temperature=1.5, top_k=50, max_tokens=5) + + tpu_llm.generate([prompt], sampling_params) + + unique_outputs = set() + for _ in range(5): + outputs = tpu_llm.generate([prompt], sampling_params) + unique_outputs.add(outputs[0].outputs[0].text.strip()) + + assert len(unique_outputs) > 1, "Random sampling failed to produce varied outputs on cached prompt." + + +def test_mixed_sampling_params_on_cached_prefix(tpu_llm: LLM): + """Verify that different sampling states are properly isolated when sharing the same cached prefix.""" + prompt_greedy = _apply_template(tpu_llm, SYSTEM_CONTEXT, "Task A: Reply with the exact word 'APPLE' and nothing else.") + prompt_random = _apply_template(tpu_llm, SYSTEM_CONTEXT, "Task B: Write a highly creative and random poem line.") + + params_greedy = SamplingParams(temperature=0, max_tokens=15) + params_random = SamplingParams(temperature=1.2, top_p=0.9, max_tokens=15) + + # Execute mixed tasks using the same prefix cache + outputs_random = tpu_llm.generate([prompt_random], params_random) + outputs_greedy = tpu_llm.generate([prompt_greedy], params_greedy) + + greedy_text = outputs_greedy[0].outputs[0].text.upper() + assert "APPLE" in greedy_text, f"Greedy task failed. Output was: '{outputs_greedy[0].outputs[0].text}'" + + +def test_logprobs_with_prefix_caching(tpu_llm: LLM): + """Verify that logprobs are correctly calculated and returned during a cache hit.""" + prompt = _apply_template(tpu_llm, SYSTEM_CONTEXT, "Explain AI in three words:") + sampling_params = SamplingParams(temperature=0, max_tokens=5, logprobs=3) + + tpu_llm.generate([prompt], sampling_params) + + outputs = tpu_llm.generate([prompt], sampling_params) + output = outputs[0].outputs[0] + + assert output.logprobs is not None, "Logprobs missing on cache hit." + for token_logprobs in output.logprobs: + for token_id, logprob_obj in token_logprobs.items(): + assert logprob_obj.logprob <= 0, f"Invalid logprob value: {logprob_obj.logprob}" \ No newline at end of file From cd143f9a5dbc384754c1892be3e4f31cf110020b Mon Sep 17 00:00:00 2001 From: Lyle Lai Date: Wed, 10 Jun 2026 09:58:59 +0000 Subject: [PATCH 37/37] fix OOM of offload_prefix_spec and chunked_offload_spec --- .../combo/test_combo_chunked_offload_spec.py | 349 ++++++++++-------- .../combo/test_combo_offload_prefix_spec.py | 323 +++++++++------- 2 files changed, 375 insertions(+), 297 deletions(-) diff --git a/tests/e2e/combo/test_combo_chunked_offload_spec.py b/tests/e2e/combo/test_combo_chunked_offload_spec.py index 6df9f2da82..6ce6bbb05b 100644 --- a/tests/e2e/combo/test_combo_chunked_offload_spec.py +++ b/tests/e2e/combo/test_combo_chunked_offload_spec.py @@ -12,191 +12,228 @@ # See the License for the specific language governing permissions and # limitations under the License. -import itertools + +import gc import os import time from typing import Optional, Union + import pytest from vllm import LLM, SamplingParams from vllm.config import KVTransferConfig + + def parse_outputs(outputs): - output_token_ids = [] - generated_texts = [] - for output in outputs: - completion = output.outputs[0] - generated_text = completion.text - token_ids = completion.token_ids - generated_texts.append(generated_text) - output_token_ids.append(token_ids) - return generated_texts, output_token_ids + output_token_ids = [] + generated_texts = [] + for output in outputs: + completion = output.outputs[0] + generated_text = completion.text + token_ids = completion.token_ids + generated_texts.append(generated_text) + output_token_ids.append(token_ids) + return generated_texts, output_token_ids + + def get_sampling_config(): - """deterministic sampling config""" - return SamplingParams(temperature=0.0, - max_tokens=20, - seed=42, - ignore_eos=True) + """deterministic sampling config""" + return SamplingParams(temperature=0.0, + max_tokens=20, + seed=42, + ignore_eos=True) + + def get_kv_transfer_config(): - """use TPUOffloadConnector""" - return KVTransferConfig( - kv_connector="TPUOffloadConnector", - kv_role="kv_both", - kv_connector_module_path="tpu_inference.offload.tpu_offload_connector", - ) + """use TPUOffloadConnector""" + return KVTransferConfig( + kv_connector="TPUOffloadConnector", + kv_role="kv_both", + kv_connector_module_path="tpu_inference.offload.tpu_offload_connector", + ) + + def _test_combo_chunked_offload_spec_accuracy( - monkeypatch: pytest.MonkeyPatch, - model_name: str, - speculative_config: Optional[Union[dict, str]], - enable_chunked_prefill: bool, - async_scheduling: Union[bool, str] = "auto", - cpu_chunks: str = "8", - max_output_len: Optional[int] = None, + monkeypatch: pytest.MonkeyPatch, + model_name: str, + speculative_config: Optional[Union[dict, str]], + enable_chunked_prefill: bool, + async_scheduling: Union[bool, str] = "auto", + cpu_chunks: str = "8", + max_output_len: Optional[int] = None, ): - sampling_config = get_sampling_config() - if max_output_len is not None: - sampling_config.max_tokens = max_output_len - kv_transfer_config = get_kv_transfer_config() - - prompts = [ - "The quick brown fox jumps over the lazy dog.", - "Artificial intelligence is transformational for scientific research and computing.", - ] - num_requests = len(prompts) - - with monkeypatch.context(): - # Standard environmental configuration for JAX-TPU KV Offloading - os.environ['SKIP_JAX_PRECOMPILE'] = '0' - os.environ['TPU_OFFLOAD_SKIP_JAX_PRECOMPILE'] = '0' - os.environ['TPU_OFFLOAD_DECODE_SAVE'] = '1' - os.environ['TPU_OFFLOAD_BATCHED_SAVE'] = '0' - os.environ['TPU_OFFLOAD_NUM_CPU_CHUNKS'] = cpu_chunks - - # Ensure Hugging Face respects offline cache if specified - os.environ["HF_HOME"] = os.environ.get("HF_HOME", "~/.cache/huggingface") - os.environ["HF_HUB_CACHE"] = os.environ.get("HF_HUB_CACHE", "~/.cache/huggingface/hub") - - tensor_parallel_size = int(os.environ.get("TPU_TP_SIZE", "8")) - - if async_scheduling == "auto": - async_scheduling_val = not speculative_config - else: - async_scheduling_val = bool(async_scheduling) - - llm = LLM( - model=model_name, - max_model_len=1024, - async_scheduling=async_scheduling_val, - tensor_parallel_size=tensor_parallel_size, - enable_chunked_prefill=enable_chunked_prefill, # Chunked Prefill Switch - kv_transfer_config=kv_transfer_config, # F03: KV Cache Offloading - speculative_config=speculative_config, # F10: Speculative Decoding - ) - - # --- Pass 1: Cold Generation (Calculates and offloads KV Cache) --- - print(f"\n--- Pass 1: Generating for {num_requests} requests ---") - t0 = time.time() - outputs1 = llm.generate(prompts, sampling_config) - pass1_time = time.time() - t0 - print(f"Pass 1 generation completed in {pass1_time:.4f} seconds") - out_texts1, out_tokens1 = parse_outputs(outputs1) - time.sleep(5) - - # --- Resetting prefix cache in TPU HBM --- - # Forces next pass to load prefix KV cache from Host CPU DRAM instead of recalculating - print("\n--- Resetting prefix cache (evicting from TPU HBM) ---") - llm.llm_engine.engine_core.reset_prefix_cache() - time.sleep(2) - - # --- Pass 2: Warm Generation (Loads KV Cache from CPU DRAM) --- - print(f"\n--- Pass 2: Generating again (should load from CPU DRAM) ---") - t0 = time.time() - outputs2 = llm.generate(prompts, sampling_config) - pass2_time = time.time() - t0 - print(f"Pass 2 generation completed in {pass2_time:.4f} seconds") - out_texts2, out_tokens2 = parse_outputs(outputs2) - time.sleep(1) - - print("\n" + "=" * 80) - print("Accuracy Comparison Results") - print(f"Pass 1 generate time: {pass1_time * 1000:.2f} ms") - print(f"Pass 2 generate time: {pass2_time * 1000:.2f} ms") - print("=" * 80) - for i in range(len(out_texts1)): - print(f"\nRequest {i}:") - print(f" Pass 1 Output Text: {out_texts1[i]!r}") - print(f" Pass 2 Output Text: {out_texts2[i]!r}") - print(f" Pass 1 Output Tokens: {out_tokens1[i]}") - print(f" Pass 2 Output Tokens: {out_tokens2[i]}") - print("\n" + "=" * 80) - - # Output 1 and Output 2 must be bit-for-bit identical - assert len(out_texts1) == len(out_texts2) - assert len(out_tokens1) == len(out_tokens2) - for i in range(len(out_texts1)): - assert out_texts1[i] == out_texts2[i], f"Text mismatch in request {i}" - assert out_tokens1[i] == out_tokens2[i], f"Token mismatch in request {i}" - - del llm - # Waiting for TPUs to be released. - time.sleep(10) + sampling_config = get_sampling_config() + if max_output_len is not None: + sampling_config.max_tokens = max_output_len + kv_transfer_config = get_kv_transfer_config() + + + prompts = [ + "The quick brown fox jumps over the lazy dog.", + "Artificial intelligence is transformational for scientific research and computing.", + ] + num_requests = len(prompts) + + + with monkeypatch.context(): + os.environ['SKIP_JAX_PRECOMPILE'] = '0' + os.environ['TPU_OFFLOAD_SKIP_JAX_PRECOMPILE'] = '1' + os.environ['TPU_OFFLOAD_DECODE_SAVE'] = '1' + os.environ['TPU_OFFLOAD_BATCHED_SAVE'] = '0' + os.environ['TPU_OFFLOAD_NUM_CPU_CHUNKS'] = cpu_chunks + + + # Ensure Hugging Face respects offline cache if specified + os.environ["HF_HOME"] = os.environ.get("HF_HOME", "~/.cache/huggingface") + os.environ["HF_HUB_CACHE"] = os.environ.get("HF_HUB_CACHE", "~/.cache/huggingface/hub") + + + tensor_parallel_size = int(os.environ.get("TPU_TP_SIZE", "8")) + + + if async_scheduling == "auto": + async_scheduling_val = not speculative_config + else: + async_scheduling_val = bool(async_scheduling) + + + llm = None + try: + llm = LLM( + model=model_name, + max_model_len=512, + max_num_batched_tokens=2048, + max_num_seqs=num_requests, + async_scheduling=async_scheduling_val, + tensor_parallel_size=tensor_parallel_size, + enable_chunked_prefill=enable_chunked_prefill, + kv_transfer_config=kv_transfer_config, + speculative_config=speculative_config, + ) + + + # --- Pass 1: Cold Generation (Calculates and offloads KV Cache) --- + print(f"\n--- Pass 1: Generating for {num_requests} requests ---") + t0 = time.time() + outputs1 = llm.generate(prompts, sampling_config) + pass1_time = time.time() - t0 + print(f"Pass 1 generation completed in {pass1_time:.4f} seconds") + out_texts1, out_tokens1 = parse_outputs(outputs1) + del outputs1 + time.sleep(5) + + + # --- Resetting prefix cache in TPU HBM --- + # Forces next pass to load prefix KV cache from Host CPU DRAM instead of recalculating + print("\n--- Resetting prefix cache (evicting from TPU HBM) ---") + llm.llm_engine.engine_core.reset_prefix_cache() + time.sleep(2) + + + # --- Pass 2: Warm Generation (Loads KV Cache from CPU DRAM) --- + print("\n--- Pass 2: Generating again (should load from CPU DRAM) ---") + t0 = time.time() + outputs2 = llm.generate(prompts, sampling_config) + pass2_time = time.time() - t0 + print(f"Pass 2 generation completed in {pass2_time:.4f} seconds") + out_texts2, out_tokens2 = parse_outputs(outputs2) + del outputs2 + time.sleep(1) + + + print("\n" + "=" * 80) + print("Accuracy Comparison Results") + print(f"Pass 1 generate time: {pass1_time * 1000:.2f} ms") + print(f"Pass 2 generate time: {pass2_time * 1000:.2f} ms") + print("=" * 80) + for i in range(len(out_texts1)): + print(f"\nRequest {i}:") + print(f" Pass 1 Output Text: {out_texts1[i]!r}") + print(f" Pass 2 Output Text: {out_texts2[i]!r}") + print(f" Pass 1 Output Tokens: {out_tokens1[i]}") + print(f" Pass 2 Output Tokens: {out_tokens2[i]}") + print("\n" + "=" * 80) + + + # Output 1 and Output 2 must be bit-for-bit identical + assert len(out_texts1) == len(out_texts2) + assert len(out_tokens1) == len(out_tokens2) + for i in range(len(out_texts1)): + assert out_texts1[i] == out_texts2[i], f"Text mismatch in request {i}" + assert out_tokens1[i] == out_tokens2[i], f"Token mismatch in request {i}" + finally: + if llm is not None and hasattr(llm.llm_engine, "shutdown"): + llm.llm_engine.shutdown() + del llm + gc.collect() + # Waiting for TPUs to be released. + time.sleep(10) + + @pytest.mark.parametrize("enable_chunked_prefill", [True, False]) @pytest.mark.parametrize("async_scheduling", [ - pytest.param(True, marks=pytest.mark.xfail(reason="Async scheduling is not supported with speculative decoding on TPU backend yet", strict=False)), - pytest.param(False) + pytest.param(True, marks=pytest.mark.skip(reason="Async scheduling is not supported with speculative decoding on TPU backend yet")), + pytest.param(False) ]) def test_combo_chunked_ngram_llama_3b(monkeypatch: pytest.MonkeyPatch, enable_chunked_prefill: bool, async_scheduling: bool): - """Tests Llama-3.2-3B-Instruct using Ngram Speculative Decoding, Chunked Prefill, and KV Offloading""" - model_name = os.environ.get("MODEL_NAME", "meta-llama/Llama-3.2-3B-Instruct") - speculative_config = { - "method": "ngram", - "prompt_lookup_max": 2, - "prompt_lookup_min": 2, - "num_speculative_tokens": 4, - } - _test_combo_chunked_offload_spec_accuracy( - monkeypatch=monkeypatch, - model_name=model_name, - speculative_config=speculative_config, - enable_chunked_prefill=enable_chunked_prefill, - async_scheduling=async_scheduling, - cpu_chunks="8", - ) + """Tests Llama-3.2-3B-Instruct using Ngram Speculative Decoding, Chunked Prefill, and KV Offloading""" + model_name = os.environ.get("MODEL_NAME", "meta-llama/Llama-3.2-3B-Instruct") + speculative_config = { + "method": "ngram", + "prompt_lookup_max": 2, + "prompt_lookup_min": 2, + "num_speculative_tokens": 4, + } + _test_combo_chunked_offload_spec_accuracy( + monkeypatch=monkeypatch, + model_name=model_name, + speculative_config=speculative_config, + enable_chunked_prefill=enable_chunked_prefill, + async_scheduling=async_scheduling, + cpu_chunks="8", + ) + + @pytest.mark.parametrize("enable_chunked_prefill", [True, False]) @pytest.mark.parametrize("async_scheduling", [ - pytest.param(True, marks=pytest.mark.xfail(reason="Async scheduling is not supported with speculative decoding on TPU backend yet", strict=False)), - pytest.param(False) + pytest.param(True, marks=pytest.mark.skip(reason="Async scheduling is not supported with speculative decoding on TPU backend yet")), + pytest.param(False) ]) def test_combo_chunked_eagle3_llama_8b(monkeypatch: pytest.MonkeyPatch, enable_chunked_prefill: bool, async_scheduling: bool): - """Tests Llama-3 8B using Eagle3 Speculative Decoding, Chunked Prefill, and KV Offloading""" - model_name = os.environ.get("MODEL_NAME", "meta-llama/Meta-Llama-3.1-8B-Instruct") - draft_model = os.environ.get("DRAFT_MODEL_NAME", "unkmaster/EAGLE3-LLaMA3.1-Instruct-8B") - - model_impl = os.environ.get("MODEL_IMPL_TYPE", "auto") - monkeypatch.setenv("DRAFT_MODEL_IMPL_TYPE", model_impl) - - speculative_config = { - "method": "eagle3", - "model": draft_model, - "num_speculative_tokens": 3, - "draft_tensor_parallel_size": 1, - } - _test_combo_chunked_offload_spec_accuracy( - monkeypatch=monkeypatch, - model_name=model_name, - speculative_config=speculative_config, - enable_chunked_prefill=enable_chunked_prefill, - async_scheduling=async_scheduling, - cpu_chunks="8", - ) + """Tests Llama-3 8B using Eagle3 Speculative Decoding, Chunked Prefill, and KV Offloading""" + model_name = os.environ.get("MODEL_NAME", "meta-llama/Meta-Llama-3.1-8B-Instruct") + draft_model = os.environ.get("DRAFT_MODEL_NAME", "unkmaster/EAGLE3-LLaMA3.1-Instruct-8B") + + model_impl = os.environ.get("MODEL_IMPL_TYPE", "auto") + monkeypatch.setenv("DRAFT_MODEL_IMPL_TYPE", model_impl) + + + speculative_config = { + "method": "eagle3", + "model": draft_model, + "num_speculative_tokens": 3, + "draft_tensor_parallel_size": 1, + } + _test_combo_chunked_offload_spec_accuracy( + monkeypatch=monkeypatch, + model_name=model_name, + speculative_config=speculative_config, + enable_chunked_prefill=enable_chunked_prefill, + async_scheduling=async_scheduling, + cpu_chunks="8", + ) + + + diff --git a/tests/e2e/combo/test_combo_offload_prefix_spec.py b/tests/e2e/combo/test_combo_offload_prefix_spec.py index 6d04e7c94e..70c01efa4d 100644 --- a/tests/e2e/combo/test_combo_offload_prefix_spec.py +++ b/tests/e2e/combo/test_combo_offload_prefix_spec.py @@ -12,170 +12,211 @@ # See the License for the specific language governing permissions and # limitations under the License. -import itertools + +import gc import os import time from typing import Optional, Union + import pytest from vllm import LLM, SamplingParams from vllm.config import KVTransferConfig + + def parse_outputs(outputs): - output_token_ids = [] - generated_texts = [] - for output in outputs: - completion = output.outputs[0] - generated_text = completion.text - token_ids = completion.token_ids - generated_texts.append(generated_text) - output_token_ids.append(token_ids) - return generated_texts, output_token_ids + output_token_ids = [] + generated_texts = [] + for output in outputs: + completion = output.outputs[0] + generated_text = completion.text + token_ids = completion.token_ids + generated_texts.append(generated_text) + output_token_ids.append(token_ids) + return generated_texts, output_token_ids + + def get_sampling_config(): - """deterministic sampling config""" - return SamplingParams(temperature=0.0, - max_tokens=20, - seed=42, - ignore_eos=True) + """deterministic sampling config""" + return SamplingParams(temperature=0.0, + max_tokens=20, + seed=42, + ignore_eos=True) + + def get_kv_transfer_config(): - """use TPUOffloadConnector""" - return KVTransferConfig( - kv_connector="TPUOffloadConnector", - kv_role="kv_both", - kv_connector_module_path="tpu_inference.offload.tpu_offload_connector", - ) + """use TPUOffloadConnector""" + return KVTransferConfig( + kv_connector="TPUOffloadConnector", + kv_role="kv_both", + kv_connector_module_path="tpu_inference.offload.tpu_offload_connector", + ) + + def _test_combo_offload_prefix_spec_accuracy( - monkeypatch: pytest.MonkeyPatch, - model_name: str, - speculative_config: Optional[Union[dict, str]], - cpu_chunks: str = "8", - max_output_len: Optional[int] = None, + monkeypatch: pytest.MonkeyPatch, + model_name: str, + speculative_config: Optional[Union[dict, str]], + cpu_chunks: str = "8", + max_output_len: Optional[int] = None, ): - sampling_config = get_sampling_config() - if max_output_len is not None: - sampling_config.max_tokens = max_output_len - kv_transfer_config = get_kv_transfer_config() - - prompts = [ - "The quick brown fox jumps over the lazy dog.", - "Artificial intelligence is transformational for scientific research and computing.", - ] - num_requests = len(prompts) - - with monkeypatch.context(): - # Standard environmental configuration for JAX-TPU KV Offloading - os.environ['SKIP_JAX_PRECOMPILE'] = '0' - os.environ['TPU_OFFLOAD_SKIP_JAX_PRECOMPILE'] = '0' - os.environ['TPU_OFFLOAD_DECODE_SAVE'] = '1' - os.environ['TPU_OFFLOAD_BATCHED_SAVE'] = '0' - os.environ['TPU_OFFLOAD_NUM_CPU_CHUNKS'] = cpu_chunks - - # Ensure Hugging Face respects offline cache if specified - os.environ["HF_HOME"] = os.environ.get("HF_HOME", "~/.cache/huggingface") - os.environ["HF_HUB_CACHE"] = os.environ.get("HF_HUB_CACHE", "~/.cache/huggingface/hub") - - tensor_parallel_size = int(os.environ.get("TPU_TP_SIZE", "8")) - - llm = LLM( - model=model_name, - max_model_len=1024, - async_scheduling=not speculative_config, - tensor_parallel_size=tensor_parallel_size, - enable_prefix_caching=True, - kv_transfer_config=kv_transfer_config, - speculative_config=speculative_config, - ) - - # --- Pass 1: Cold Generation (Calculates and offloads KV Cache) --- - print(f"\n--- Pass 1: Generating for {num_requests} requests ---") - t0 = time.time() - outputs1 = llm.generate(prompts, sampling_config) - pass1_time = time.time() - t0 - print(f"Pass 1 generation completed in {pass1_time:.4f} seconds") - out_texts1, out_tokens1 = parse_outputs(outputs1) - time.sleep(5) - - # --- Resetting prefix cache in TPU HBM --- - # Forces next pass to load prefix KV cache from Host CPU DRAM instead of recalculating - print("\n--- Resetting prefix cache (evicting from TPU HBM) ---") - llm.llm_engine.engine_core.reset_prefix_cache() - time.sleep(2) - - # --- Pass 2: Warm Generation (Loads KV Cache from CPU DRAM) --- - print(f"\n--- Pass 2: Generating again (should load from CPU DRAM) ---") - t0 = time.time() - outputs2 = llm.generate(prompts, sampling_config) - pass2_time = time.time() - t0 - print(f"Pass 2 generation completed in {pass2_time:.4f} seconds") - out_texts2, out_tokens2 = parse_outputs(outputs2) - time.sleep(1) - - print("\n" + "=" * 80) - print("Accuracy Comparison Results") - print(f"Pass 1 generate time: {pass1_time * 1000:.2f} ms") - print(f"Pass 2 generate time: {pass2_time * 1000:.2f} ms") - print("=" * 80) - for i in range(len(out_texts1)): - print(f"\nRequest {i}:") - print(f" Pass 1 Output Text: {out_texts1[i]!r}") - print(f" Pass 2 Output Text: {out_texts2[i]!r}") - print(f" Pass 1 Output Tokens: {out_tokens1[i]}") - print(f" Pass 2 Output Tokens: {out_tokens2[i]}") - print("\n" + "=" * 80) - - # Output 1 and Output 2 must be bit-for-bit identical - assert len(out_texts1) == len(out_texts2) - assert len(out_tokens1) == len(out_tokens2) - for i in range(len(out_texts1)): - assert out_texts1[i] == out_texts2[i], f"Text mismatch in request {i}" - assert out_tokens1[i] == out_tokens2[i], f"Token mismatch in request {i}" - - del llm - # Waiting for TPUs to be released. - time.sleep(10) + sampling_config = get_sampling_config() + if max_output_len is not None: + sampling_config.max_tokens = max_output_len + kv_transfer_config = get_kv_transfer_config() + + + prompts = [ + "The quick brown fox jumps over the lazy dog.", + "Artificial intelligence is transformational for scientific research and computing.", + ] + num_requests = len(prompts) + + + with monkeypatch.context(): + # Standard environmental configuration for JAX-TPU KV Offloading + monkeypatch.setenv('SKIP_JAX_PRECOMPILE', '0') + monkeypatch.setenv('TPU_OFFLOAD_SKIP_JAX_PRECOMPILE', '1') + monkeypatch.setenv('TPU_OFFLOAD_DECODE_SAVE', '1') + monkeypatch.setenv('TPU_OFFLOAD_BATCHED_SAVE', '0') + monkeypatch.setenv('TPU_OFFLOAD_NUM_CPU_CHUNKS', cpu_chunks) + + + # Ensure Hugging Face respects offline cache if specified + monkeypatch.setenv("HF_HOME", + os.environ.get("HF_HOME", "~/.cache/huggingface")) + monkeypatch.setenv( + "HF_HUB_CACHE", + os.environ.get("HF_HUB_CACHE", "~/.cache/huggingface/hub")) + + + tensor_parallel_size = int(os.environ.get("TPU_TP_SIZE", "8")) + + + llm = None + try: + llm = LLM( + model=model_name, + max_model_len=512, + max_num_batched_tokens=2048, + max_num_seqs=num_requests, + async_scheduling=not speculative_config, + tensor_parallel_size=tensor_parallel_size, + enable_prefix_caching=True, + kv_transfer_config=kv_transfer_config, + speculative_config=speculative_config, + ) + + + # --- Pass 1: Cold Generation (Calculates and offloads KV Cache) --- + print(f"\n--- Pass 1: Generating for {num_requests} requests ---") + t0 = time.time() + outputs1 = llm.generate(prompts, sampling_config) + pass1_time = time.time() - t0 + print(f"Pass 1 generation completed in {pass1_time:.4f} seconds") + out_texts1, out_tokens1 = parse_outputs(outputs1) + del outputs1 + time.sleep(5) + + + # --- Resetting prefix cache in TPU HBM --- + # Forces next pass to load prefix KV cache from Host CPU DRAM instead of recalculating + print("\n--- Resetting prefix cache (evicting from TPU HBM) ---") + llm.llm_engine.engine_core.reset_prefix_cache() + time.sleep(2) + + + # --- Pass 2: Warm Generation (Loads KV Cache from CPU DRAM) --- + print("\n--- Pass 2: Generating again (should load from CPU DRAM) ---") + t0 = time.time() + outputs2 = llm.generate(prompts, sampling_config) + pass2_time = time.time() - t0 + print(f"Pass 2 generation completed in {pass2_time:.4f} seconds") + out_texts2, out_tokens2 = parse_outputs(outputs2) + del outputs2 + time.sleep(1) + + + print("\n" + "=" * 80) + print("Accuracy Comparison Results") + print(f"Pass 1 generate time: {pass1_time * 1000:.2f} ms") + print(f"Pass 2 generate time: {pass2_time * 1000:.2f} ms") + print("=" * 80) + for i in range(len(out_texts1)): + print(f"\nRequest {i}:") + print(f" Pass 1 Output Text: {out_texts1[i]!r}") + print(f" Pass 2 Output Text: {out_texts2[i]!r}") + print(f" Pass 1 Output Tokens: {out_tokens1[i]}") + print(f" Pass 2 Output Tokens: {out_tokens2[i]}") + print("\n" + "=" * 80) + + + # Output 1 and Output 2 must be bit-for-bit identical + assert len(out_texts1) == len(out_texts2) + assert len(out_tokens1) == len(out_tokens2) + for i in range(len(out_texts1)): + assert out_texts1[i] == out_texts2[i], f"Text mismatch in request {i}" + assert out_tokens1[i] == out_tokens2[i], f"Token mismatch in request {i}" + finally: + if llm is not None and hasattr(llm.llm_engine, "shutdown"): + llm.llm_engine.shutdown() + del llm + gc.collect() + # Waiting for TPUs to be released. + time.sleep(10) + + def test_combo_ngram_llama_3b(monkeypatch: pytest.MonkeyPatch): - """Tests Llama-3.2-3B-Instruct using Ngram Speculative Decoding, Prefix Caching, and KV Offloading""" - model_name = os.environ.get("MODEL_NAME", "meta-llama/Llama-3.2-3B-Instruct") - speculative_config = { - "method": "ngram", - "prompt_lookup_max": 2, - "prompt_lookup_min": 2, - "num_speculative_tokens": 4, - } - _test_combo_offload_prefix_spec_accuracy( - monkeypatch=monkeypatch, - model_name=model_name, - speculative_config=speculative_config, - cpu_chunks="8", - ) + """Tests Llama-3.2-3B-Instruct using Ngram Speculative Decoding, Prefix Caching, and KV Offloading""" + model_name = os.environ.get("MODEL_NAME", "meta-llama/Llama-3.2-3B-Instruct") + speculative_config = { + "method": "ngram", + "prompt_lookup_max": 2, + "prompt_lookup_min": 2, + "num_speculative_tokens": 4, + } + _test_combo_offload_prefix_spec_accuracy( + monkeypatch=monkeypatch, + model_name=model_name, + speculative_config=speculative_config, + cpu_chunks="8", + ) + + def test_combo_eagle3_llama_8b(monkeypatch: pytest.MonkeyPatch): - """Tests Llama-3 8B using Eagle3 Speculative Decoding, Prefix Caching, and KV Offloading""" - model_name = os.environ.get("MODEL_NAME", "meta-llama/Meta-Llama-3.1-8B-Instruct") - draft_model = os.environ.get("DRAFT_MODEL_NAME", "unkmaster/EAGLE3-LLaMA3.1-Instruct-8B") - - model_impl = os.environ.get("MODEL_IMPL_TYPE", "auto") - monkeypatch.setenv("DRAFT_MODEL_IMPL_TYPE", model_impl) - - speculative_config = { - "method": "eagle3", - "model": draft_model, - "num_speculative_tokens": 3, - "draft_tensor_parallel_size": 1, - } - _test_combo_offload_prefix_spec_accuracy( - monkeypatch=monkeypatch, - model_name=model_name, - speculative_config=speculative_config, - cpu_chunks="8", - ) + """Tests Llama-3 8B using Eagle3 Speculative Decoding, Prefix Caching, and KV Offloading""" + model_name = os.environ.get("MODEL_NAME", "meta-llama/Meta-Llama-3.1-8B-Instruct") + draft_model = os.environ.get("DRAFT_MODEL_NAME", "unkmaster/EAGLE3-LLaMA3.1-Instruct-8B") + + + model_impl = os.environ.get("MODEL_IMPL_TYPE", "auto") + monkeypatch.setenv("DRAFT_MODEL_IMPL_TYPE", model_impl) + + + speculative_config = { + "method": "eagle3", + "model": draft_model, + "num_speculative_tokens": 3, + "draft_tensor_parallel_size": 1, + } + _test_combo_offload_prefix_spec_accuracy( + monkeypatch=monkeypatch, + model_name=model_name, + speculative_config=speculative_config, + cpu_chunks="8", + ) + + +