From b65b338850bc1291ed290f99d000ff18e312e57d Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Sun, 19 Jul 2026 18:21:16 +0800 Subject: [PATCH 1/8] perf(musa): add chunked Qwen min-p sampler --- csrc/musa/min_p_sampler.mu | 386 ++++++++++++++++++++++++++++ csrc/musa/musa_ops.h | 6 + csrc/musa/torch_bindings.cpp | 7 + setup.py | 1 + tests/test_chunked_min_p_sampler.py | 167 ++++++++++++ vllm_musa/_custom_ops.py | 36 ++- vllm_musa/utils/environ.py | 1 + 7 files changed, 603 insertions(+), 1 deletion(-) create mode 100644 csrc/musa/min_p_sampler.mu create mode 100644 tests/test_chunked_min_p_sampler.py diff --git a/csrc/musa/min_p_sampler.mu b/csrc/musa/min_p_sampler.mu new file mode 100644 index 000000000000..ea87a725e073 --- /dev/null +++ b/csrc/musa/min_p_sampler.mu @@ -0,0 +1,386 @@ +/* + * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * + * Shape-specialized min-p sampling for Qwen's 151936-token vocabulary. + * The chunked reduction and multi-CTA dataflow are adapted from RubyMine + * problem 900015, submission 20732. The production wrapper below restores + * FlashInfer's Philox, per-row threshold, and row-index contracts. + */ + +#include +#include +#include +#include + +#include +#include +#include + +#include "musa.h" +#include "pytorch_extension_utils.h" +#include "torch_musa/csrc/aten/musa/MUSAGeneratorImpl.h" +#include "torch_musa/csrc/core/MUSAGuard.h" +#include "torch_musa/csrc/core/MUSAStream.h" + +namespace { + +constexpr int kThreads = 256; +constexpr int kWarp = 32; +constexpr int kRowsPerBlock = kThreads / kWarp; +constexpr int kChunk = 2048; +constexpr int kQwenVocab = 151936; +constexpr int kMaxBatch = 64; + +struct alignas(16) Float4 { + float x, y, z, w; +}; + +__device__ __forceinline__ Float4 load4_cached(const float* ptr) { + return *reinterpret_cast(ptr); +} + +__device__ __forceinline__ Float4 load4_bypass(const float* ptr) { + // The standalone RubyMine donor uses ld.global.cs.v4.f32. In the full + // vLLM extension, combining that four-output inline assembly with Philox + // exhausts the mp31 output-register allocator. Keep the chunked multi-CTA + // dataflow first; a later isolated ticket can reintroduce cache bypass with + // a lower-pressure load helper if end-to-end evidence warrants it. + return load4_cached(ptr); +} + +__device__ __forceinline__ float warp_max(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value = fmaxf(value, __shfl_down_sync(0xffffffffu, value, offset)); + } + return value; +} + +__device__ __forceinline__ float warp_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset); + } + return value; +} + +__device__ __forceinline__ float block_max(float value) { + __shared__ float warp_values[8]; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + value = warp_max(value); + if (lane == 0) { + warp_values[warp] = value; + } + __syncthreads(); + value = threadIdx.x < 8 ? warp_values[lane] : -FLT_MAX; + if (warp == 0) { + value = warp_max(value); + } + return __shfl_sync(0xffffffffu, value, 0); +} + +__device__ __forceinline__ float block_sum(float value) { + __shared__ float warp_values[8]; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + value = warp_sum(value); + if (lane == 0) { + warp_values[warp] = value; + } + __syncthreads(); + value = threadIdx.x < 8 ? warp_values[lane] : 0.0f; + if (warp == 0) { + value = warp_sum(value); + } + return __shfl_sync(0xffffffffu, value, 0); +} + +__global__ void partial_max_kernel(const float* __restrict__ probs, + const int* __restrict__ indices, + float* __restrict__ partial, int vocab, + int chunks) { + const int job = blockIdx.x; + const int row = job / chunks; + const int chunk = job - row * chunks; + const int row_idx = indices == nullptr ? row : indices[row]; + const int begin = chunk * kChunk; + const int end = min(vocab, begin + kChunk); + const float* row_probs = probs + static_cast(row_idx) * vocab; + float local = -FLT_MAX; + for (int index = begin + threadIdx.x * 4; index < end; + index += blockDim.x * 4) { + if (index + 3 < end) { + const Float4 value = load4_bypass(row_probs + index); + local = fmaxf(local, value.x); + local = fmaxf(local, value.y); + local = fmaxf(local, value.z); + local = fmaxf(local, value.w); + } else { + for (int tail = index; tail < end; ++tail) { + local = fmaxf(local, row_probs[tail]); + } + } + } + const float maximum = block_max(local); + if (threadIdx.x == 0) { + partial[job] = maximum; + } +} + +__global__ void threshold_rng_kernel( + const float* __restrict__ partial, float* __restrict__ threshold, + float* __restrict__ uniform, const float* __restrict__ min_p_arr, + float min_p_val, int* __restrict__ output, int batch, int vocab, + int chunks, uint64_t philox_seed, uint64_t philox_offset) { + const int warp = threadIdx.x / kWarp; + const int lane = threadIdx.x & (kWarp - 1); + const int row = blockIdx.x * kRowsPerBlock + warp; + if (row >= batch) { + return; + } + + float maximum = -FLT_MAX; + for (int chunk = lane; chunk < chunks; chunk += kWarp) { + maximum = fmaxf(maximum, partial[row * chunks + chunk]); + } + maximum = warp_max(maximum); + if (lane == 0) { + const float min_p = min_p_arr == nullptr ? min_p_val : min_p_arr[row]; + murandStatePhilox4_32_10_t state; + murand_init(philox_seed, row, philox_offset, &state); + threshold[row] = maximum * min_p; + uniform[row] = murand_uniform(&state); + output[row] = vocab - 1; + } +} + +__global__ void partial_sum_kernel(const float* __restrict__ probs, + const int* __restrict__ indices, + const float* __restrict__ threshold, + float* __restrict__ partial, int vocab, + int chunks) { + const int job = blockIdx.x; + const int row = job / chunks; + const int chunk = job - row * chunks; + const int row_idx = indices == nullptr ? row : indices[row]; + const int begin = chunk * kChunk; + const int end = min(vocab, begin + kChunk); + const float* row_probs = probs + static_cast(row_idx) * vocab; + const float cutoff = threshold[row]; + float local = 0.0f; + for (int index = begin + threadIdx.x * 4; index < end; + index += blockDim.x * 4) { + if (index + 3 < end) { + const Float4 value = load4_bypass(row_probs + index); + local += value.x >= cutoff ? value.x : 0.0f; + local += value.y >= cutoff ? value.y : 0.0f; + local += value.z >= cutoff ? value.z : 0.0f; + local += value.w >= cutoff ? value.w : 0.0f; + } else { + for (int tail = index; tail < end; ++tail) { + const float value = row_probs[tail]; + local += value >= cutoff ? value : 0.0f; + } + } + } + const float total = block_sum(local); + if (threadIdx.x == 0) { + partial[job] = total; + } +} + +__device__ __forceinline__ float warp_prefix(float value) { + const int lane = threadIdx.x & 31; +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const float incoming = __shfl_up_sync(0xffffffffu, value, offset); + if (lane >= offset) { + value += incoming; + } + } + return value; +} + +__global__ void select_kernel(const float* __restrict__ probs, + const int* __restrict__ indices, + const float* __restrict__ threshold, + const float* __restrict__ uniform, + const float* __restrict__ partial, + int* __restrict__ output, int batch, int vocab, + int chunks) { + const int warp = threadIdx.x / kWarp; + const int lane = threadIdx.x & (kWarp - 1); + const int row = blockIdx.x * kRowsPerBlock + warp; + if (row >= batch) { + return; + } + + __shared__ int selected_chunks[kRowsPerBlock]; + __shared__ float selected_targets[kRowsPerBlock]; + __shared__ int selected_tokens[kRowsPerBlock]; + __shared__ int last_valid_tokens[kRowsPerBlock]; + + if (lane == 0) { + float total = 0.0f; + int last_nonzero_chunk = 0; + for (int chunk = 0; chunk < chunks; ++chunk) { + const float chunk_sum = partial[row * chunks + chunk]; + total += chunk_sum; + if (chunk_sum > 0.0f) { + last_nonzero_chunk = chunk; + } + } + const float target = uniform[row] * total; + float prefix = 0.0f; + int selected = last_nonzero_chunk; + for (int chunk = 0; chunk < chunks; ++chunk) { + const float next = prefix + partial[row * chunks + chunk]; + if (next > target) { + selected = chunk; + break; + } + prefix = next; + } + selected_chunks[warp] = selected; + selected_targets[warp] = target - prefix; + selected_tokens[warp] = vocab; + last_valid_tokens[warp] = -1; + } + __syncwarp(); + + const int chunk = selected_chunks[warp]; + const int begin = chunk * kChunk; + const int end = min(vocab, begin + kChunk); + const int row_idx = indices == nullptr ? row : indices[row]; + const float* row_probs = probs + static_cast(row_idx) * vocab; + const float cutoff = threshold[row]; + const int interval = kChunk / kWarp; + const int lane_begin = begin + lane * interval; + const int lane_end = min(end, lane_begin + interval); + + float lane_sum = 0.0f; + int lane_last_valid = -1; + for (int index = lane_begin; index < lane_end; ++index) { + const float value = row_probs[index]; + if (value >= cutoff) { + lane_sum += value; + lane_last_valid = index; + } + } + const float inclusive = warp_prefix(lane_sum); + const float lane_prefix = inclusive - lane_sum; + if (lane_last_valid >= 0) { + atomicMax(&last_valid_tokens[warp], lane_last_valid); + } + + const float target = selected_targets[warp]; + if (lane_prefix <= target && inclusive > target) { + float cumulative = lane_prefix; + for (int index = lane_begin; index < lane_end; ++index) { + const float value = row_probs[index]; + if (value >= cutoff) { + cumulative += value; + if (cumulative > target) { + atomicMin(&selected_tokens[warp], index); + break; + } + } + } + } + __syncwarp(); + if (lane == 0) { + const int selected = selected_tokens[warp]; + const int fallback = last_valid_tokens[warp] >= 0 + ? last_valid_tokens[warp] + : vocab - 1; + output[row] = selected < vocab ? selected : fallback; + } +} + +} // namespace + +void musa_chunked_min_p_sampling_from_probs( + at::Tensor probs, at::Tensor output, + std::optional maybe_indices, + std::optional maybe_min_p_arr, double min_p_val, + bool deterministic, std::optional gen_) { + CHECK_INPUT(probs); + CHECK_INPUT(output); + CHECK_DIM(2, probs); + CHECK_DIM(1, output); + CHECK_EQ(probs.dtype(), torch::kFloat); + CHECK_EQ(output.dtype(), torch::kInt); + CHECK_EQ(probs.device(), output.device()); + + const int batch = output.size(0); + const int vocab = probs.size(1); + TORCH_CHECK(batch > 0 && batch <= kMaxBatch, + "chunked min-p sampler requires 1 <= batch <= ", kMaxBatch); + TORCH_CHECK(vocab == kQwenVocab, + "chunked min-p sampler requires vocab=", kQwenVocab); + TORCH_CHECK(probs.size(0) >= batch, + "chunked min-p sampler received fewer probability rows than outputs"); + + int* indices = nullptr; + if (maybe_indices.has_value()) { + at::Tensor indices_tensor = maybe_indices.value(); + CHECK_INPUT(indices_tensor); + CHECK_EQ(indices_tensor.dtype(), torch::kInt); + CHECK_EQ(indices_tensor.device(), probs.device()); + TORCH_CHECK(indices_tensor.numel() >= batch, + "chunked min-p indices must cover every output row"); + indices = indices_tensor.data_ptr(); + } + + float* min_p_arr = nullptr; + if (maybe_min_p_arr.has_value()) { + at::Tensor min_p_tensor = maybe_min_p_arr.value(); + CHECK_INPUT(min_p_tensor); + CHECK_EQ(min_p_tensor.dtype(), torch::kFloat); + CHECK_EQ(min_p_tensor.device(), probs.device()); + TORCH_CHECK(min_p_tensor.numel() >= batch, + "chunked min-p thresholds must cover every output row"); + min_p_arr = min_p_tensor.data_ptr(); + } + + auto gen = at::get_generator_or_default( + gen_, at::musa::detail::getDefaultMUSAGenerator()); + uint64_t philox_seed; + uint64_t philox_offset; + { + std::lock_guard lock(gen->mutex_); + // Match FlashInfer's production min-p contract exactly: one Philox + // counter reservation per sampled row. + at::PhiloxMusaState state = gen->philox_musa_state(batch); + philox_seed = state.seed_.val; + philox_offset = state.offset_.val; + } + + const int chunks = (vocab + kChunk - 1) / kChunk; + auto scratch = at::empty({static_cast(batch) * (chunks + 2)}, + probs.options()); + float* partial = scratch.data_ptr(); + float* threshold = partial + static_cast(batch) * chunks; + float* uniform = threshold + batch; + + const c10::musa::OptionalMUSAGuard device_guard(probs.device()); + auto stream = at::musa::getCurrentMUSAStream(); + partial_max_kernel<<>>( + probs.data_ptr(), indices, partial, vocab, chunks); + threshold_rng_kernel<<<(batch + kRowsPerBlock - 1) / kRowsPerBlock, + kThreads, 0, stream>>>( + partial, threshold, uniform, min_p_arr, static_cast(min_p_val), + output.data_ptr(), batch, vocab, chunks, philox_seed, philox_offset); + partial_sum_kernel<<>>( + probs.data_ptr(), indices, threshold, partial, vocab, chunks); + select_kernel<<<(batch + kRowsPerBlock - 1) / kRowsPerBlock, kThreads, 0, + stream>>>(probs.data_ptr(), indices, threshold, uniform, + partial, output.data_ptr(), batch, vocab, + chunks); + const musaError_t error = musaGetLastError(); + TORCH_CHECK(error == musaSuccess, "chunked min-p sampler failed: ", + musaGetErrorString(error)); + (void)deterministic; +} diff --git a/csrc/musa/musa_ops.h b/csrc/musa/musa_ops.h index 14a9cc93d193..7d91a1dffae5 100644 --- a/csrc/musa/musa_ops.h +++ b/csrc/musa/musa_ops.h @@ -72,6 +72,12 @@ void musa_top_k_top_p_sampling_from_probs( bool deterministic, std::optional gen_); +void musa_chunked_min_p_sampling_from_probs( + at::Tensor probs, at::Tensor output, + std::optional maybe_indices, + std::optional maybe_min_p_arr, double min_p_val, + bool deterministic, std::optional gen_); + /* * From FlashInfer */ diff --git a/csrc/musa/torch_bindings.cpp b/csrc/musa/torch_bindings.cpp index ddb594f9d7d7..229f6b518da6 100644 --- a/csrc/musa/torch_bindings.cpp +++ b/csrc/musa/torch_bindings.cpp @@ -60,6 +60,13 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _musa_ops), musa_ops) { "float top_k_val, Tensor? maybe_top_p_arr, float top_p_val, bool deterministic, Generator? gen) -> ()"); musa_ops.impl("musa_top_k_top_p_sampling_from_probs", torch::kMUSA, &musa_top_k_top_p_sampling_from_probs); + musa_ops.def( + "musa_chunked_min_p_sampling_from_probs(Tensor probs, Tensor! output, " + "Tensor? maybe_indices, Tensor? maybe_min_p_arr, float min_p_val, " + "bool deterministic, Generator? gen) -> ()"); + musa_ops.impl("musa_chunked_min_p_sampling_from_probs", torch::kMUSA, + &musa_chunked_min_p_sampling_from_probs); + /* * From FlashInfer */ diff --git a/setup.py b/setup.py index 39a9932b5a72..787e9d75d939 100644 --- a/setup.py +++ b/setup.py @@ -221,6 +221,7 @@ def __init__(self, name, git_repository, git_tag, git_shallow=False): "csrc/musa/quantization/fused_add_rms_norm_per_token_group_fp8_quant.cu", "csrc/musa/quantization/per_token_group_quant_8bit_vec.cu", "csrc/musa/sampler.mu", + "csrc/musa/min_p_sampler.mu", str(_FLASHINFER_REPO.source_dir / "csrc/norm.cu"), str(_FLASHINFER_REPO.source_dir / "csrc/renorm.cu"), str(_FLASHINFER_REPO.source_dir / "csrc/sampling.cu"), diff --git a/tests/test_chunked_min_p_sampler.py b/tests/test_chunked_min_p_sampler.py new file mode 100644 index 000000000000..3d7bb414dd02 --- /dev/null +++ b/tests/test_chunked_min_p_sampler.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +"""S5000 contract checks for the chunked Qwen min-p sampler.""" + +from __future__ import annotations + +from collections.abc import Iterator + +import pytest + +pytest.importorskip("torchada") +torch = pytest.importorskip("torch") +pytest.importorskip("torch_musa") + +VOCAB = 151936 + + +@pytest.fixture(scope="module", autouse=True) +def _musa_device() -> Iterator[None]: + if not hasattr(torch, "musa") or not torch.musa.is_available(): + pytest.skip("MUSA device is not available") + torch.musa.set_device(0) + + from vllm.platforms import current_platform + + import vllm_musa + + if not current_platform.is_device_capability((3, 1)): + pytest.skip("the chunked min-p sampler requires S5000/mp31") + vllm_musa.register_custom_ops() + yield + + +def _probs(batch: int) -> torch.Tensor: + generator = torch.Generator(device="cpu").manual_seed(763000 + batch) + logits = torch.randn((batch, VOCAB), generator=generator, dtype=torch.float32) + return torch.softmax(logits.mul_(1.5), dim=-1).contiguous().to("musa") + + +def _candidate( + probs: torch.Tensor, + min_p: float | torch.Tensor, + *, + indices: torch.Tensor | None = None, + seed: int = 763, +) -> torch.Tensor: + output = torch.empty(probs.size(0), dtype=torch.int32, device="musa") + generator = torch.Generator(device="musa").manual_seed(seed) + min_p_arr = min_p if isinstance(min_p, torch.Tensor) else None + min_p_val = 0.0 if min_p_arr is not None else min_p + torch.ops._C_musa_ops.musa_chunked_min_p_sampling_from_probs( + probs, + output, + indices, + min_p_arr, + min_p_val, + True, + generator, + ) + return output + + +def _assert_support( + probs: torch.Tensor, + samples: torch.Tensor, + min_p: float | torch.Tensor, + indices: torch.Tensor | None = None, +) -> None: + rows = ( + torch.arange(probs.size(0), dtype=torch.int32, device="musa") + if indices is None + else indices + ) + row_probs = probs[rows.long()] + chosen = row_probs.gather(1, samples.long()[:, None]).view(-1) + threshold = row_probs.amax(dim=1) * min_p + assert bool((chosen + 1.0e-12 >= threshold).all().item()) + + +@pytest.mark.parametrize("batch", [1, 4, 16, 64]) +def test_seeded_candidate_is_repeatable_and_in_support(batch: int) -> None: + probs = _probs(batch) + first = _candidate(probs, 0.05, seed=1000 + batch) + second = _candidate(probs, 0.05, seed=1000 + batch) + torch.musa.synchronize() + + assert first.dtype == torch.int32 + assert first.shape == (batch,) + assert torch.equal(first, second) + assert bool(((first >= 0) & (first < VOCAB)).all().item()) + _assert_support(probs, first, 0.05) + + +def test_per_row_threshold_and_row_indices_contract() -> None: + batch = 16 + probs = _probs(batch) + indices = torch.arange(batch - 1, -1, -1, dtype=torch.int32, device="musa") + min_p = torch.linspace(0.01, 0.20, batch, dtype=torch.float32, device="musa") + first = _candidate(probs, min_p, indices=indices, seed=9001) + second = _candidate(probs, min_p, indices=indices, seed=9001) + torch.musa.synchronize() + + assert torch.equal(first, second) + _assert_support(probs, first, min_p, indices) + + +@pytest.mark.parametrize("batch", [1, 4, 16, 64]) +def test_candidate_matches_production_philox_contract(batch: int) -> None: + probs = _probs(batch) + current_output = torch.empty(batch, dtype=torch.int32, device="musa") + candidate_output = torch.empty(batch, dtype=torch.int32, device="musa") + current_generator = torch.Generator(device="musa").manual_seed(763100 + batch) + candidate_generator = torch.Generator(device="musa").manual_seed(763100 + batch) + + torch.ops._C_musa_ops.min_p_sampling_from_probs( + probs, current_output, None, None, 0.05, True, current_generator + ) + torch.ops._C_musa_ops.musa_chunked_min_p_sampling_from_probs( + probs, candidate_output, None, None, 0.05, True, candidate_generator + ) + assert torch.equal(current_generator.get_state(), candidate_generator.get_state()) + current_next = torch.rand(8, generator=current_generator, device="musa") + candidate_next = torch.rand(8, generator=candidate_generator, device="musa") + torch.musa.synchronize() + + assert torch.equal(current_output, candidate_output) + assert torch.equal(current_next, candidate_next) + + +@pytest.mark.parametrize("min_p", [0.0, 1.0]) +def test_min_p_edge_values(min_p: float) -> None: + probs = _probs(4) + samples = _candidate(probs, min_p, seed=1234) + torch.musa.synchronize() + _assert_support(probs, samples, min_p) + if min_p == 1.0: + selected = probs.gather(1, samples.long()[:, None]).view(-1) + torch.testing.assert_close(selected, probs.amax(dim=1), rtol=0, atol=0) + + +def test_custom_op_wrapper_falls_back_outside_qwen_shape(monkeypatch) -> None: + from vllm_musa import _custom_ops + from vllm_musa.utils.environ import envs + + probs = torch.softmax(torch.randn((4, 1024), device="musa"), dim=-1) + with envs.VLLM_MUSA_CHUNKED_MIN_P_SAMPLER.override(True): + samples = _custom_ops.min_p_sampling_from_probs(probs, 0.05) + torch.musa.synchronize() + + assert samples.dtype == torch.int32 + assert samples.shape == (4,) + assert not envs.VLLM_MUSA_CHUNKED_MIN_P_SAMPLER.is_set() + _assert_support(probs, samples, 0.05) + + +def test_candidate_captures_and_replays() -> None: + probs = _probs(4) + output = torch.empty(4, dtype=torch.int32, device="musa") + graph = torch.musa.MUSAGraph() + with torch.musa.graph(graph): + torch.ops._C_musa_ops.musa_chunked_min_p_sampling_from_probs( + probs, output, None, None, 0.05, True, None + ) + graph.replay() + torch.musa.synchronize() + + assert bool(((output >= 0) & (output < VOCAB)).all().item()) + _assert_support(probs, output, 0.05) diff --git a/vllm_musa/_custom_ops.py b/vllm_musa/_custom_ops.py index d72bf71c2b08..386c15363f30 100644 --- a/vllm_musa/_custom_ops.py +++ b/vllm_musa/_custom_ops.py @@ -3,6 +3,8 @@ import torch +from vllm_musa.utils.environ import envs + try: import vllm_musa._C # noqa: F401 except ImportError as e: @@ -326,7 +328,39 @@ def _min_p_sampling_from_probs_internal( probs = probs.float() maybe_min_p_arr = maybe_min_p_arr.float() if maybe_min_p_arr is not None else None samples = torch.empty(probs.size(0), dtype=torch.int32, device=device) - torch.ops._C_musa_ops.min_p_sampling_from_probs.default( + use_chunked = ( + envs.VLLM_MUSA_CHUNKED_MIN_P_SAMPLER.get() + and deterministic + and probs.ndim == 2 + and 0 < probs.shape[0] <= 64 + and probs.shape[1] == 151936 + and probs.device.type == "musa" + and probs.is_contiguous() + and ( + indices is None + or ( + indices.dtype == torch.int32 + and indices.device == probs.device + and indices.is_contiguous() + and indices.numel() >= probs.shape[0] + ) + ) + and ( + maybe_min_p_arr is None + or ( + maybe_min_p_arr.device == probs.device + and maybe_min_p_arr.dtype == torch.float32 + and maybe_min_p_arr.is_contiguous() + and maybe_min_p_arr.numel() >= probs.shape[0] + ) + ) + ) + op = ( + torch.ops._C_musa_ops.musa_chunked_min_p_sampling_from_probs.default + if use_chunked + else torch.ops._C_musa_ops.min_p_sampling_from_probs.default + ) + op( probs, samples, indices, diff --git a/vllm_musa/utils/environ.py b/vllm_musa/utils/environ.py index 3e50734ddba7..1041d20629ce 100644 --- a/vllm_musa/utils/environ.py +++ b/vllm_musa/utils/environ.py @@ -91,6 +91,7 @@ class Envs: VLLM_MUSA_FUSED_ADD_RMSNORM = EnvBool(True) VLLM_MUSA_ENABLE_JIT_TOPK = EnvBool(True) VLLM_MUSA_SAMPLER_FAST_PATH = EnvBool(True) + VLLM_MUSA_CHUNKED_MIN_P_SAMPLER = EnvBool(False) VLLM_MUSA_SEEDED_MULTINOMIAL = EnvBool(True) VLLM_MUSA_RESHAPE_CACHE_FLASH = EnvBool(True) From cb02740d0ecad7649a232cd58dbfe901d5958aab Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Sun, 19 Jul 2026 19:20:01 +0800 Subject: [PATCH 2/8] perf(musa): cover Qwen3.5 sampling vocab --- csrc/musa/min_p_sampler.mu | 10 +++++---- tests/test_chunked_min_p_sampler.py | 33 ++++++++++++++++++----------- vllm_musa/_custom_ops.py | 2 +- 3 files changed, 28 insertions(+), 17 deletions(-) diff --git a/csrc/musa/min_p_sampler.mu b/csrc/musa/min_p_sampler.mu index ea87a725e073..bcf63d79849c 100644 --- a/csrc/musa/min_p_sampler.mu +++ b/csrc/musa/min_p_sampler.mu @@ -2,7 +2,7 @@ * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd. * SPDX-License-Identifier: Apache-2.0 * - * Shape-specialized min-p sampling for Qwen's 151936-token vocabulary. + * Shape-specialized min-p sampling for Qwen's production vocabularies. * The chunked reduction and multi-CTA dataflow are adapted from RubyMine * problem 900015, submission 20732. The production wrapper below restores * FlashInfer's Philox, per-row threshold, and row-index contracts. @@ -29,7 +29,8 @@ constexpr int kThreads = 256; constexpr int kWarp = 32; constexpr int kRowsPerBlock = kThreads / kWarp; constexpr int kChunk = 2048; -constexpr int kQwenVocab = 151936; +constexpr int kQwen3Vocab = 151936; +constexpr int kQwen35Vocab = 248320; constexpr int kMaxBatch = 64; struct alignas(16) Float4 { @@ -318,8 +319,9 @@ void musa_chunked_min_p_sampling_from_probs( const int vocab = probs.size(1); TORCH_CHECK(batch > 0 && batch <= kMaxBatch, "chunked min-p sampler requires 1 <= batch <= ", kMaxBatch); - TORCH_CHECK(vocab == kQwenVocab, - "chunked min-p sampler requires vocab=", kQwenVocab); + TORCH_CHECK(vocab == kQwen3Vocab || vocab == kQwen35Vocab, + "chunked min-p sampler requires vocab=", kQwen3Vocab, " or ", + kQwen35Vocab); TORCH_CHECK(probs.size(0) >= batch, "chunked min-p sampler received fewer probability rows than outputs"); diff --git a/tests/test_chunked_min_p_sampler.py b/tests/test_chunked_min_p_sampler.py index 3d7bb414dd02..b9cace4e0f7c 100644 --- a/tests/test_chunked_min_p_sampler.py +++ b/tests/test_chunked_min_p_sampler.py @@ -11,7 +11,9 @@ torch = pytest.importorskip("torch") pytest.importorskip("torch_musa") -VOCAB = 151936 +QWEN3_VOCAB = 151936 +QWEN35_VOCAB = 248320 +SUPPORTED_VOCABS = (QWEN3_VOCAB, QWEN35_VOCAB) @pytest.fixture(scope="module", autouse=True) @@ -30,9 +32,9 @@ def _musa_device() -> Iterator[None]: yield -def _probs(batch: int) -> torch.Tensor: - generator = torch.Generator(device="cpu").manual_seed(763000 + batch) - logits = torch.randn((batch, VOCAB), generator=generator, dtype=torch.float32) +def _probs(batch: int, vocab: int = QWEN3_VOCAB) -> torch.Tensor: + generator = torch.Generator(device="cpu").manual_seed(763000 + batch + vocab) + logits = torch.randn((batch, vocab), generator=generator, dtype=torch.float32) return torch.softmax(logits.mul_(1.5), dim=-1).contiguous().to("musa") @@ -76,9 +78,12 @@ def _assert_support( assert bool((chosen + 1.0e-12 >= threshold).all().item()) +@pytest.mark.parametrize("vocab", SUPPORTED_VOCABS) @pytest.mark.parametrize("batch", [1, 4, 16, 64]) -def test_seeded_candidate_is_repeatable_and_in_support(batch: int) -> None: - probs = _probs(batch) +def test_seeded_candidate_is_repeatable_and_in_support( + batch: int, vocab: int +) -> None: + probs = _probs(batch, vocab) first = _candidate(probs, 0.05, seed=1000 + batch) second = _candidate(probs, 0.05, seed=1000 + batch) torch.musa.synchronize() @@ -86,7 +91,7 @@ def test_seeded_candidate_is_repeatable_and_in_support(batch: int) -> None: assert first.dtype == torch.int32 assert first.shape == (batch,) assert torch.equal(first, second) - assert bool(((first >= 0) & (first < VOCAB)).all().item()) + assert bool(((first >= 0) & (first < vocab)).all().item()) _assert_support(probs, first, 0.05) @@ -103,9 +108,12 @@ def test_per_row_threshold_and_row_indices_contract() -> None: _assert_support(probs, first, min_p, indices) +@pytest.mark.parametrize("vocab", SUPPORTED_VOCABS) @pytest.mark.parametrize("batch", [1, 4, 16, 64]) -def test_candidate_matches_production_philox_contract(batch: int) -> None: - probs = _probs(batch) +def test_candidate_matches_production_philox_contract( + batch: int, vocab: int +) -> None: + probs = _probs(batch, vocab) current_output = torch.empty(batch, dtype=torch.int32, device="musa") candidate_output = torch.empty(batch, dtype=torch.int32, device="musa") current_generator = torch.Generator(device="musa").manual_seed(763100 + batch) @@ -152,8 +160,9 @@ def test_custom_op_wrapper_falls_back_outside_qwen_shape(monkeypatch) -> None: _assert_support(probs, samples, 0.05) -def test_candidate_captures_and_replays() -> None: - probs = _probs(4) +@pytest.mark.parametrize("vocab", SUPPORTED_VOCABS) +def test_candidate_captures_and_replays(vocab: int) -> None: + probs = _probs(4, vocab) output = torch.empty(4, dtype=torch.int32, device="musa") graph = torch.musa.MUSAGraph() with torch.musa.graph(graph): @@ -163,5 +172,5 @@ def test_candidate_captures_and_replays() -> None: graph.replay() torch.musa.synchronize() - assert bool(((output >= 0) & (output < VOCAB)).all().item()) + assert bool(((output >= 0) & (output < probs.shape[1])).all().item()) _assert_support(probs, output, 0.05) diff --git a/vllm_musa/_custom_ops.py b/vllm_musa/_custom_ops.py index 386c15363f30..7e748ec5599e 100644 --- a/vllm_musa/_custom_ops.py +++ b/vllm_musa/_custom_ops.py @@ -333,7 +333,7 @@ def _min_p_sampling_from_probs_internal( and deterministic and probs.ndim == 2 and 0 < probs.shape[0] <= 64 - and probs.shape[1] == 151936 + and probs.shape[1] in (151936, 248320) and probs.device.type == "musa" and probs.is_contiguous() and ( From e22535cc7ef05484543051e92f685b92dab550c4 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Mon, 20 Jul 2026 11:10:38 +0800 Subject: [PATCH 3/8] perf(musa): enable guarded Qwen min-p sampler by default --- tests/test_chunked_min_p_sampler.py | 86 ++++++++++++-- tests/test_chunked_min_p_sampler_source.py | 45 ++++++++ vllm_musa/_custom_ops.py | 128 ++++++++++++++++----- vllm_musa/utils/environ.py | 1 - 4 files changed, 221 insertions(+), 39 deletions(-) create mode 100644 tests/test_chunked_min_p_sampler_source.py diff --git a/tests/test_chunked_min_p_sampler.py b/tests/test_chunked_min_p_sampler.py index b9cace4e0f7c..9e905061afc5 100644 --- a/tests/test_chunked_min_p_sampler.py +++ b/tests/test_chunked_min_p_sampler.py @@ -80,9 +80,7 @@ def _assert_support( @pytest.mark.parametrize("vocab", SUPPORTED_VOCABS) @pytest.mark.parametrize("batch", [1, 4, 16, 64]) -def test_seeded_candidate_is_repeatable_and_in_support( - batch: int, vocab: int -) -> None: +def test_seeded_candidate_is_repeatable_and_in_support(batch: int, vocab: int) -> None: probs = _probs(batch, vocab) first = _candidate(probs, 0.05, seed=1000 + batch) second = _candidate(probs, 0.05, seed=1000 + batch) @@ -110,9 +108,7 @@ def test_per_row_threshold_and_row_indices_contract() -> None: @pytest.mark.parametrize("vocab", SUPPORTED_VOCABS) @pytest.mark.parametrize("batch", [1, 4, 16, 64]) -def test_candidate_matches_production_philox_contract( - batch: int, vocab: int -) -> None: +def test_candidate_matches_production_philox_contract(batch: int, vocab: int) -> None: probs = _probs(batch, vocab) current_output = torch.empty(batch, dtype=torch.int32, device="musa") candidate_output = torch.empty(batch, dtype=torch.int32, device="musa") @@ -145,21 +141,89 @@ def test_min_p_edge_values(min_p: float) -> None: torch.testing.assert_close(selected, probs.amax(dim=1), rtol=0, atol=0) -def test_custom_op_wrapper_falls_back_outside_qwen_shape(monkeypatch) -> None: +def test_custom_op_wrapper_falls_back_outside_qwen_shape() -> None: from vllm_musa import _custom_ops - from vllm_musa.utils.environ import envs probs = torch.softmax(torch.randn((4, 1024), device="musa"), dim=-1) - with envs.VLLM_MUSA_CHUNKED_MIN_P_SAMPLER.override(True): - samples = _custom_ops.min_p_sampling_from_probs(probs, 0.05) + samples = _custom_ops.min_p_sampling_from_probs(probs, 0.05) torch.musa.synchronize() assert samples.dtype == torch.int32 assert samples.shape == (4,) - assert not envs.VLLM_MUSA_CHUNKED_MIN_P_SAMPLER.is_set() _assert_support(probs, samples, 0.05) +def test_default_dispatch_contract() -> None: + from vllm_musa import _custom_ops + + probs = _probs(1) + assert _custom_ops._can_use_chunked_min_p_sampler(probs, None, None, True, None) + + +def test_default_wrapper_preserves_philox_contract() -> None: + from vllm_musa import _custom_ops + + probs = _probs(1) + expected = torch.empty(1, dtype=torch.int32, device="musa") + expected_generator = torch.Generator(device="musa").manual_seed(763200) + torch.ops._C_musa_ops.min_p_sampling_from_probs( + probs, expected, None, None, 0.05, True, expected_generator + ) + + actual_generator = torch.Generator(device="musa").manual_seed(763200) + actual = _custom_ops.min_p_sampling_from_probs( + probs, 0.05, generator=actual_generator + ) + torch.musa.synchronize() + + assert torch.equal(actual, expected) + assert torch.equal(actual_generator.get_state(), expected_generator.get_state()) + + +def test_unsupported_contracts_and_seeded_cpu_fallback() -> None: + from vllm_musa import _custom_ops + + probs = _probs(1) + musa_generator = torch.Generator(device="musa") + cpu_generator = torch.Generator(device="cpu") + + assert _custom_ops._can_use_chunked_min_p_sampler( + probs, None, None, True, musa_generator + ) + assert not _custom_ops._can_use_chunked_min_p_sampler( + probs, None, None, True, cpu_generator + ) + assert not _custom_ops._can_use_chunked_min_p_sampler( + probs, None, None, False, None + ) + assert not _custom_ops._can_use_chunked_min_p_sampler( + probs.half(), None, None, True, None + ) + assert not _custom_ops._can_use_chunked_min_p_sampler( + _probs(1, 1024), None, None, True, None + ) + assert not _custom_ops._can_use_chunked_min_p_sampler( + probs.expand(65, -1), None, None, True, None + ) + + non_contiguous = torch.empty( + (1, 2 * probs.shape[1]), dtype=torch.float32, device="musa" + )[:, ::2] + assert not non_contiguous.is_contiguous() + assert not _custom_ops._can_use_chunked_min_p_sampler( + non_contiguous, None, None, True, None + ) + + bad_indices = torch.arange(1, dtype=torch.int64, device="musa") + assert not _custom_ops._can_use_chunked_min_p_sampler( + probs, bad_indices, None, True, None + ) + bad_min_p = torch.full((1,), 0.05, dtype=torch.float16, device="musa") + assert not _custom_ops._can_use_chunked_min_p_sampler( + probs, None, bad_min_p, True, None + ) + + @pytest.mark.parametrize("vocab", SUPPORTED_VOCABS) def test_candidate_captures_and_replays(vocab: int) -> None: probs = _probs(4, vocab) diff --git a/tests/test_chunked_min_p_sampler_source.py b/tests/test_chunked_min_p_sampler_source.py new file mode 100644 index 000000000000..736f61db8a85 --- /dev/null +++ b/tests/test_chunked_min_p_sampler_source.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Source contracts for the default-on RubyMine min-p dispatch.""" + +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] + + +def _read(relative_path: str) -> str: + return (ROOT / relative_path).read_text(encoding="utf-8") + + +def test_min_p_dispatch_has_no_environment_gate() -> None: + source = _read("vllm_musa/_custom_ops.py") + environ = _read("vllm_musa/utils/environ.py") + + assert "envs." not in source + assert "CHUNKED_MIN_P_SAMPLER" not in environ + + +def test_min_p_dispatch_keeps_the_validated_contract_guard() -> None: + source = _read("vllm_musa/_custom_ops.py") + + assert "def _can_use_chunked_min_p_sampler(" in source + assert "deterministic" in source + assert "and probs.dtype == torch.float32" in source + assert "and probs.shape[1] in _QWEN_MIN_P_VOCABS" in source + assert "and 0 < probs.shape[0] <= _QWEN_MIN_P_MAX_BATCH" in source + assert "and probs.is_contiguous()" in source + assert "and _is_validated_musa_device(probs.device)" in source + assert "device_id=device_id" in source + assert "and _is_supported_musa_generator(generator, probs.device)" in source + assert "musa_chunked_min_p_sampling_from_probs.default" in source + assert "min_p_sampling_from_probs.default" in source + + +def test_min_p_dispatch_preserves_unsupported_input_fallbacks() -> None: + source = _read("vllm_musa/_custom_ops.py") + + assert "indices.dtype == torch.int32" in source + assert "indices.is_contiguous()" in source + assert "maybe_min_p_arr.dtype == torch.float32" in source + assert "maybe_min_p_arr.is_contiguous()" in source + assert "input_probs = probs" in source + assert "input_min_p_arr = maybe_min_p_arr" in source diff --git a/vllm_musa/_custom_ops.py b/vllm_musa/_custom_ops.py index 7e748ec5599e..29bf1b194139 100644 --- a/vllm_musa/_custom_ops.py +++ b/vllm_musa/_custom_ops.py @@ -3,7 +3,8 @@ import torch -from vllm_musa.utils.environ import envs +_QWEN_MIN_P_VOCABS = (151936, 248320) +_QWEN_MIN_P_MAX_BATCH = 64 try: import vllm_musa._C # noqa: F401 @@ -172,6 +173,97 @@ def _to_tensor_scalar_tuple(x): return (None, x) +def _musa_device_index(device: torch.device) -> Optional[int]: + """Resolve a logical MUSA device index without changing the current device.""" + try: + device_id = device.index + if device_id is None: + device_id = int(torch.musa.current_device()) + if device_id < 0 or device_id >= torch.musa.device_count(): + return None + return device_id + except Exception: + return None + + +def _is_validated_musa_device(device: torch.device) -> bool: + """Return whether the tensor's device is the validated S5000 target.""" + try: + from vllm.platforms import current_platform + except Exception: + return False + + try: + device_id = _musa_device_index(device) + if device_id is None: + return False + return current_platform.is_device_capability((3, 1), device_id=device_id) + except Exception: + # A device query can fail while the platform is still initializing. + # Falling back is safer than selecting an architecture-specialized op. + return False + + +def _is_supported_musa_generator( + generator: Optional[torch.Generator], device: torch.device +) -> bool: + """Check that a seeded generator belongs to the sampled MUSA device.""" + if generator is None: + return True + try: + generator_device = generator.device + if getattr(generator_device, "type", None) != "musa": + return False + tensor_device_id = _musa_device_index(device) + generator_device_id = _musa_device_index(generator_device) + return ( + tensor_device_id is not None + and generator_device_id is not None + and tensor_device_id == generator_device_id + ) + except Exception: + return False + + +def _can_use_chunked_min_p_sampler( + probs: torch.Tensor, + indices: Optional[torch.Tensor], + maybe_min_p_arr: Optional[torch.Tensor], + deterministic: bool, + generator: Optional[torch.Generator], +) -> bool: + """Gate the RubyMine min-p kernel to its validated production contract.""" + return ( + deterministic + and probs.device.type == "musa" + and _is_validated_musa_device(probs.device) + and probs.dtype == torch.float32 + and probs.ndim == 2 + and 0 < probs.shape[0] <= _QWEN_MIN_P_MAX_BATCH + and probs.shape[1] in _QWEN_MIN_P_VOCABS + and probs.is_contiguous() + and _is_supported_musa_generator(generator, probs.device) + and ( + indices is None + or ( + indices.dtype == torch.int32 + and indices.device == probs.device + and indices.is_contiguous() + and indices.numel() >= probs.shape[0] + ) + ) + and ( + maybe_min_p_arr is None + or ( + maybe_min_p_arr.device == probs.device + and maybe_min_p_arr.dtype == torch.float32 + and maybe_min_p_arr.is_contiguous() + and maybe_min_p_arr.numel() >= probs.shape[0] + ) + ) + ) + + def _top_k_renorm_probs_internal( probs: torch.Tensor, maybe_top_k_arr: Optional[torch.Tensor], @@ -325,35 +417,17 @@ def _min_p_sampling_from_probs_internal( generator: Optional[torch.Generator], ) -> torch.Tensor: device = probs.device + input_probs = probs + input_min_p_arr = maybe_min_p_arr probs = probs.float() maybe_min_p_arr = maybe_min_p_arr.float() if maybe_min_p_arr is not None else None samples = torch.empty(probs.size(0), dtype=torch.int32, device=device) - use_chunked = ( - envs.VLLM_MUSA_CHUNKED_MIN_P_SAMPLER.get() - and deterministic - and probs.ndim == 2 - and 0 < probs.shape[0] <= 64 - and probs.shape[1] in (151936, 248320) - and probs.device.type == "musa" - and probs.is_contiguous() - and ( - indices is None - or ( - indices.dtype == torch.int32 - and indices.device == probs.device - and indices.is_contiguous() - and indices.numel() >= probs.shape[0] - ) - ) - and ( - maybe_min_p_arr is None - or ( - maybe_min_p_arr.device == probs.device - and maybe_min_p_arr.dtype == torch.float32 - and maybe_min_p_arr.is_contiguous() - and maybe_min_p_arr.numel() >= probs.shape[0] - ) - ) + use_chunked = _can_use_chunked_min_p_sampler( + input_probs, + indices, + input_min_p_arr, + deterministic, + generator, ) op = ( torch.ops._C_musa_ops.musa_chunked_min_p_sampling_from_probs.default diff --git a/vllm_musa/utils/environ.py b/vllm_musa/utils/environ.py index 1041d20629ce..3e50734ddba7 100644 --- a/vllm_musa/utils/environ.py +++ b/vllm_musa/utils/environ.py @@ -91,7 +91,6 @@ class Envs: VLLM_MUSA_FUSED_ADD_RMSNORM = EnvBool(True) VLLM_MUSA_ENABLE_JIT_TOPK = EnvBool(True) VLLM_MUSA_SAMPLER_FAST_PATH = EnvBool(True) - VLLM_MUSA_CHUNKED_MIN_P_SAMPLER = EnvBool(False) VLLM_MUSA_SEEDED_MULTINOMIAL = EnvBool(True) VLLM_MUSA_RESHAPE_CACHE_FLASH = EnvBool(True) From 6dc7f9fcfcfff0c3e00ff102ecd1254526dbbfa4 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Mon, 20 Jul 2026 12:21:52 +0800 Subject: [PATCH 4/8] style(musa): keep min-p comments code-local --- csrc/musa/min_p_sampler.mu | 12 ++++-------- tests/test_chunked_min_p_sampler_source.py | 2 +- vllm_musa/_custom_ops.py | 2 +- 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/csrc/musa/min_p_sampler.mu b/csrc/musa/min_p_sampler.mu index bcf63d79849c..f5fb9f8ce8f8 100644 --- a/csrc/musa/min_p_sampler.mu +++ b/csrc/musa/min_p_sampler.mu @@ -3,9 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 * * Shape-specialized min-p sampling for Qwen's production vocabularies. - * The chunked reduction and multi-CTA dataflow are adapted from RubyMine - * problem 900015, submission 20732. The production wrapper below restores - * FlashInfer's Philox, per-row threshold, and row-index contracts. + * The chunked reduction preserves FlashInfer's Philox, per-row threshold, and + * row-index contracts while splitting wide rows across multiple CTAs. */ #include @@ -42,11 +41,8 @@ __device__ __forceinline__ Float4 load4_cached(const float* ptr) { } __device__ __forceinline__ Float4 load4_bypass(const float* ptr) { - // The standalone RubyMine donor uses ld.global.cs.v4.f32. In the full - // vLLM extension, combining that four-output inline assembly with Philox - // exhausts the mp31 output-register allocator. Keep the chunked multi-CTA - // dataflow first; a later isolated ticket can reintroduce cache bypass with - // a lower-pressure load helper if end-to-end evidence warrants it. + // The cached vector load avoids the register pressure of a four-output + // inline-assembly load when combined with Philox state. return load4_cached(ptr); } diff --git a/tests/test_chunked_min_p_sampler_source.py b/tests/test_chunked_min_p_sampler_source.py index 736f61db8a85..b76ddd0443ea 100644 --- a/tests/test_chunked_min_p_sampler_source.py +++ b/tests/test_chunked_min_p_sampler_source.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Source contracts for the default-on RubyMine min-p dispatch.""" +"""Source contracts for the default-on chunked min-p dispatch.""" from pathlib import Path diff --git a/vllm_musa/_custom_ops.py b/vllm_musa/_custom_ops.py index 29bf1b194139..4b7e2033337d 100644 --- a/vllm_musa/_custom_ops.py +++ b/vllm_musa/_custom_ops.py @@ -232,7 +232,7 @@ def _can_use_chunked_min_p_sampler( deterministic: bool, generator: Optional[torch.Generator], ) -> bool: - """Gate the RubyMine min-p kernel to its validated production contract.""" + """Gate the chunked min-p kernel to its validated production contract.""" return ( deterministic and probs.device.type == "musa" From 2b6d877cf06428a7300fb19c63845863445f1489 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Mon, 20 Jul 2026 12:35:17 +0800 Subject: [PATCH 5/8] fix(musa): map visible devices for capability guards --- tests/test_chunked_min_p_sampler_source.py | 2 ++ vllm_musa/_custom_ops.py | 7 +++++++ 2 files changed, 9 insertions(+) diff --git a/tests/test_chunked_min_p_sampler_source.py b/tests/test_chunked_min_p_sampler_source.py index b76ddd0443ea..73ce1b8215f8 100644 --- a/tests/test_chunked_min_p_sampler_source.py +++ b/tests/test_chunked_min_p_sampler_source.py @@ -28,6 +28,8 @@ def test_min_p_dispatch_keeps_the_validated_contract_guard() -> None: assert "and 0 < probs.shape[0] <= _QWEN_MIN_P_MAX_BATCH" in source assert "and probs.is_contiguous()" in source assert "and _is_validated_musa_device(probs.device)" in source + assert "MtmlMUSAPlatform" in source + assert "visible_device_id_to_physical_device_id" in source assert "device_id=device_id" in source assert "and _is_supported_musa_generator(generator, probs.device)" in source assert "musa_chunked_min_p_sampling_from_probs.default" in source diff --git a/vllm_musa/_custom_ops.py b/vllm_musa/_custom_ops.py index 4b7e2033337d..f86c1ef4d8fe 100644 --- a/vllm_musa/_custom_ops.py +++ b/vllm_musa/_custom_ops.py @@ -190,6 +190,8 @@ def _is_validated_musa_device(device: torch.device) -> bool: """Return whether the tensor's device is the validated S5000 target.""" try: from vllm.platforms import current_platform + + from vllm_musa.platform import MtmlMUSAPlatform except Exception: return False @@ -197,6 +199,11 @@ def _is_validated_musa_device(device: torch.device) -> bool: device_id = _musa_device_index(device) if device_id is None: return False + if isinstance(current_platform, MtmlMUSAPlatform): + # MTML queries physical IDs while tensor device indices are visible ordinals. + device_id = current_platform.visible_device_id_to_physical_device_id( + device_id + ) return current_platform.is_device_capability((3, 1), device_id=device_id) except Exception: # A device query can fail while the platform is still initializing. From 549d2d41c6115f9dd2c65fe8b69e45598b634c71 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Mon, 20 Jul 2026 23:28:17 +0800 Subject: [PATCH 6/8] fix(musa): query logical sampler device capability --- tests/test_chunked_min_p_sampler_source.py | 71 +++++++++++++++++++++- vllm_musa/_custom_ops.py | 22 +++---- 2 files changed, 77 insertions(+), 16 deletions(-) diff --git a/tests/test_chunked_min_p_sampler_source.py b/tests/test_chunked_min_p_sampler_source.py index 73ce1b8215f8..5e3df060cfaa 100644 --- a/tests/test_chunked_min_p_sampler_source.py +++ b/tests/test_chunked_min_p_sampler_source.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 """Source contracts for the default-on chunked min-p dispatch.""" +import ast +from functools import cache from pathlib import Path +from types import SimpleNamespace +from typing import Optional ROOT = Path(__file__).resolve().parents[1] @@ -10,6 +14,31 @@ def _read(relative_path: str) -> str: return (ROOT / relative_path).read_text(encoding="utf-8") +def _load_device_guard(musa): + source_path = ROOT / "vllm_musa/_custom_ops.py" + tree = ast.parse(source_path.read_text(encoding="utf-8"), source_path) + helper_names = { + "_musa_device_index", + "_get_musa_device_capability", + "_is_validated_musa_device", + } + helpers = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name in helper_names + ] + namespace = { + "Optional": Optional, + "cache": cache, + "torch": SimpleNamespace(device=object, musa=musa), + } + exec( + compile(ast.Module(body=helpers, type_ignores=[]), source_path, "exec"), + namespace, + ) + return namespace["_is_validated_musa_device"] + + def test_min_p_dispatch_has_no_environment_gate() -> None: source = _read("vllm_musa/_custom_ops.py") environ = _read("vllm_musa/utils/environ.py") @@ -28,14 +57,50 @@ def test_min_p_dispatch_keeps_the_validated_contract_guard() -> None: assert "and 0 < probs.shape[0] <= _QWEN_MIN_P_MAX_BATCH" in source assert "and probs.is_contiguous()" in source assert "and _is_validated_musa_device(probs.device)" in source - assert "MtmlMUSAPlatform" in source - assert "visible_device_id_to_physical_device_id" in source - assert "device_id=device_id" in source + assert "_get_musa_device_capability(device_id)" in source + assert "torch.musa.get_device_capability(device_id)" in source assert "and _is_supported_musa_generator(generator, probs.device)" in source assert "musa_chunked_min_p_sampling_from_probs.default" in source assert "min_p_sampling_from_probs.default" in source +def test_device_guard_queries_tensor_logical_device_id() -> None: + queried_device_ids = [] + musa = SimpleNamespace( + current_device=lambda: 0, + device_count=lambda: 2, + get_device_capability=lambda device_id: ( + queried_device_ids.append(device_id) or (3, 1) + ), + ) + is_validated = _load_device_guard(musa) + + assert is_validated(SimpleNamespace(type="musa", index=1)) + assert is_validated(SimpleNamespace(type="musa", index=1)) + assert queried_device_ids == [1] + + +def test_device_guard_does_not_cache_runtime_query_failures() -> None: + queried_device_ids = [] + + def fail_capability_query(device_id: int) -> tuple[int, int]: + queried_device_ids.append(device_id) + if len(queried_device_ids) == 1: + raise RuntimeError("capability query failed") + return (3, 1) + + musa = SimpleNamespace( + current_device=lambda: 0, + device_count=lambda: 2, + get_device_capability=fail_capability_query, + ) + is_validated = _load_device_guard(musa) + + assert not is_validated(SimpleNamespace(type="musa", index=1)) + assert is_validated(SimpleNamespace(type="musa", index=1)) + assert queried_device_ids == [1, 1] + + def test_min_p_dispatch_preserves_unsupported_input_fallbacks() -> None: source = _read("vllm_musa/_custom_ops.py") diff --git a/vllm_musa/_custom_ops.py b/vllm_musa/_custom_ops.py index f86c1ef4d8fe..dca85659bc2a 100644 --- a/vllm_musa/_custom_ops.py +++ b/vllm_musa/_custom_ops.py @@ -1,4 +1,5 @@ import logging +from functools import cache from typing import Optional, Union import torch @@ -186,25 +187,20 @@ def _musa_device_index(device: torch.device) -> Optional[int]: return None -def _is_validated_musa_device(device: torch.device) -> bool: - """Return whether the tensor's device is the validated S5000 target.""" - try: - from vllm.platforms import current_platform +@cache +def _get_musa_device_capability(device_id: int) -> tuple[int, int]: + """Query one logical MUSA device and cache successful results.""" + return torch.musa.get_device_capability(device_id) - from vllm_musa.platform import MtmlMUSAPlatform - except Exception: - return False +def _is_validated_musa_device(device: torch.device) -> bool: + """Return whether the tensor's device is the validated S5000 target.""" try: device_id = _musa_device_index(device) if device_id is None: return False - if isinstance(current_platform, MtmlMUSAPlatform): - # MTML queries physical IDs while tensor device indices are visible ordinals. - device_id = current_platform.visible_device_id_to_physical_device_id( - device_id - ) - return current_platform.is_device_capability((3, 1), device_id=device_id) + # torch.musa accepts visible logical IDs, matching torch.device.index. + return _get_musa_device_capability(device_id) == (3, 1) except Exception: # A device query can fail while the platform is still initializing. # Falling back is safer than selecting an architecture-specialized op. From 273fa4ab81292dcb746073ec9d5eaaa9412ae4fc Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Tue, 21 Jul 2026 01:33:01 +0800 Subject: [PATCH 7/8] fix(musa): unpack Philox state inside min-p kernel --- csrc/musa/min_p_sampler.mu | 24 +++++++++++----------- tests/test_chunked_min_p_sampler.py | 18 +++++++++++----- tests/test_chunked_min_p_sampler_source.py | 14 +++++++++++++ 3 files changed, 39 insertions(+), 17 deletions(-) diff --git a/csrc/musa/min_p_sampler.mu b/csrc/musa/min_p_sampler.mu index f5fb9f8ce8f8..8e5d1c3aad60 100644 --- a/csrc/musa/min_p_sampler.mu +++ b/csrc/musa/min_p_sampler.mu @@ -15,10 +15,11 @@ #include #include #include +#include #include "musa.h" #include "pytorch_extension_utils.h" -#include "torch_musa/csrc/aten/musa/MUSAGeneratorImpl.h" +#include "torch_musa/csrc/aten/musa/MUSAGraphsUtils.muh" #include "torch_musa/csrc/core/MUSAGuard.h" #include "torch_musa/csrc/core/MUSAStream.h" @@ -130,7 +131,10 @@ __global__ void threshold_rng_kernel( const float* __restrict__ partial, float* __restrict__ threshold, float* __restrict__ uniform, const float* __restrict__ min_p_arr, float min_p_val, int* __restrict__ output, int batch, int vocab, - int chunks, uint64_t philox_seed, uint64_t philox_offset) { + int chunks, at::PhiloxMusaState philox_state) { + const auto seeds = at::musa::philox::unpack(philox_state); + const uint64_t philox_seed = std::get<0>(seeds); + const uint64_t philox_offset = std::get<1>(seeds); const int warp = threadIdx.x / kWarp; const int lane = threadIdx.x & (kWarp - 1); const int row = blockIdx.x * kRowsPerBlock + warp; @@ -343,18 +347,15 @@ void musa_chunked_min_p_sampling_from_probs( min_p_arr = min_p_tensor.data_ptr(); } + const c10::musa::OptionalMUSAGuard device_guard(probs.device()); auto gen = at::get_generator_or_default( - gen_, at::musa::detail::getDefaultMUSAGenerator()); - uint64_t philox_seed; - uint64_t philox_offset; - { + gen_, at::musa::detail::getDefaultMUSAGenerator(probs.get_device())); + const at::PhiloxMusaState philox_state = [&] { std::lock_guard lock(gen->mutex_); // Match FlashInfer's production min-p contract exactly: one Philox // counter reservation per sampled row. - at::PhiloxMusaState state = gen->philox_musa_state(batch); - philox_seed = state.seed_.val; - philox_offset = state.offset_.val; - } + return gen->philox_musa_state(batch); + }(); const int chunks = (vocab + kChunk - 1) / kChunk; auto scratch = at::empty({static_cast(batch) * (chunks + 2)}, @@ -363,14 +364,13 @@ void musa_chunked_min_p_sampling_from_probs( float* threshold = partial + static_cast(batch) * chunks; float* uniform = threshold + batch; - const c10::musa::OptionalMUSAGuard device_guard(probs.device()); auto stream = at::musa::getCurrentMUSAStream(); partial_max_kernel<<>>( probs.data_ptr(), indices, partial, vocab, chunks); threshold_rng_kernel<<<(batch + kRowsPerBlock - 1) / kRowsPerBlock, kThreads, 0, stream>>>( partial, threshold, uniform, min_p_arr, static_cast(min_p_val), - output.data_ptr(), batch, vocab, chunks, philox_seed, philox_offset); + output.data_ptr(), batch, vocab, chunks, philox_state); partial_sum_kernel<<>>( probs.data_ptr(), indices, threshold, partial, vocab, chunks); select_kernel<<<(batch + kRowsPerBlock - 1) / kRowsPerBlock, kThreads, 0, diff --git a/tests/test_chunked_min_p_sampler.py b/tests/test_chunked_min_p_sampler.py index 9e905061afc5..f1f360862d17 100644 --- a/tests/test_chunked_min_p_sampler.py +++ b/tests/test_chunked_min_p_sampler.py @@ -226,15 +226,23 @@ def test_unsupported_contracts_and_seeded_cpu_fallback() -> None: @pytest.mark.parametrize("vocab", SUPPORTED_VOCABS) def test_candidate_captures_and_replays(vocab: int) -> None: - probs = _probs(4, vocab) - output = torch.empty(4, dtype=torch.int32, device="musa") + batch = 64 + probs = torch.full((batch, vocab), 1.0 / vocab, dtype=torch.float32, device="musa") + output = torch.empty(batch, dtype=torch.int32, device="musa") + torch.musa.manual_seed(763300 + vocab) graph = torch.musa.MUSAGraph() with torch.musa.graph(graph): torch.ops._C_musa_ops.musa_chunked_min_p_sampling_from_probs( probs, output, None, None, 0.05, True, None ) - graph.replay() + replayed = [] + for _ in range(3): + graph.replay() + replayed.append(output.clone()) torch.musa.synchronize() - assert bool(((output >= 0) & (output < probs.shape[1])).all().item()) - _assert_support(probs, output, 0.05) + for samples in replayed: + assert bool(((samples >= 0) & (samples < probs.shape[1])).all().item()) + _assert_support(probs, samples, 0.05) + assert not torch.equal(replayed[0], replayed[1]) + assert not torch.equal(replayed[1], replayed[2]) diff --git a/tests/test_chunked_min_p_sampler_source.py b/tests/test_chunked_min_p_sampler_source.py index 5e3df060cfaa..40f17570e744 100644 --- a/tests/test_chunked_min_p_sampler_source.py +++ b/tests/test_chunked_min_p_sampler_source.py @@ -110,3 +110,17 @@ def test_min_p_dispatch_preserves_unsupported_input_fallbacks() -> None: assert "maybe_min_p_arr.is_contiguous()" in source assert "input_probs = probs" in source assert "input_min_p_arr = maybe_min_p_arr" in source + + +def test_min_p_kernel_uses_graph_safe_philox_state() -> None: + source = _read("csrc/musa/min_p_sampler.mu") + + assert "MUSAGraphsUtils.muh" in source + assert "at::PhiloxMusaState philox_state" in source + assert "at::musa::philox::unpack(philox_state)" in source + assert "philox_state.seed_.val" not in source + assert "philox_state.offset_.val" not in source + assert source.index("OptionalMUSAGuard device_guard") < source.index( + "getDefaultMUSAGenerator" + ) + assert "getDefaultMUSAGenerator(probs.get_device())" in source From c8ade1f9f92189ea18368912165f4d9fbb79a219 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Tue, 21 Jul 2026 01:56:30 +0800 Subject: [PATCH 8/8] fix(musa): make sampling fast path environment independent --- tests/test_chunked_min_p_sampler_source.py | 4 ++++ vllm_musa/utils/environ.py | 1 - vllm_musa/v1/sample/topk_topp_sampler.py | 26 ++++++---------------- 3 files changed, 11 insertions(+), 20 deletions(-) diff --git a/tests/test_chunked_min_p_sampler_source.py b/tests/test_chunked_min_p_sampler_source.py index 40f17570e744..2850a2eabaca 100644 --- a/tests/test_chunked_min_p_sampler_source.py +++ b/tests/test_chunked_min_p_sampler_source.py @@ -41,10 +41,14 @@ def _load_device_guard(musa): def test_min_p_dispatch_has_no_environment_gate() -> None: source = _read("vllm_musa/_custom_ops.py") + sampler = _read("vllm_musa/v1/sample/topk_topp_sampler.py") environ = _read("vllm_musa/utils/environ.py") assert "envs." not in source assert "CHUNKED_MIN_P_SAMPLER" not in environ + assert "VLLM_MUSA_SAMPLER_FAST_PATH" not in sampler + assert "VLLM_MUSA_SAMPLER_FAST_PATH" not in environ + assert "sampler_fast_path_enabled" not in sampler def test_min_p_dispatch_keeps_the_validated_contract_guard() -> None: diff --git a/vllm_musa/utils/environ.py b/vllm_musa/utils/environ.py index 3e50734ddba7..726cf5f10c09 100644 --- a/vllm_musa/utils/environ.py +++ b/vllm_musa/utils/environ.py @@ -90,7 +90,6 @@ class Envs: VLLM_MUSA_CUSTOM_OP_USE_NATIVE = EnvBool(False) VLLM_MUSA_FUSED_ADD_RMSNORM = EnvBool(True) VLLM_MUSA_ENABLE_JIT_TOPK = EnvBool(True) - VLLM_MUSA_SAMPLER_FAST_PATH = EnvBool(True) VLLM_MUSA_SEEDED_MULTINOMIAL = EnvBool(True) VLLM_MUSA_RESHAPE_CACHE_FLASH = EnvBool(True) diff --git a/vllm_musa/v1/sample/topk_topp_sampler.py b/vllm_musa/v1/sample/topk_topp_sampler.py index b7d622079b25..6f6522217073 100644 --- a/vllm_musa/v1/sample/topk_topp_sampler.py +++ b/vllm_musa/v1/sample/topk_topp_sampler.py @@ -21,10 +21,6 @@ _SAMPLING_EPS = 1e-5 -def sampler_fast_path_enabled() -> bool: - return envs.VLLM_MUSA_SAMPLER_FAST_PATH.get() - - def musa_seeded_multinomial_enabled() -> bool: return envs.VLLM_MUSA_SEEDED_MULTINOMIAL.get() @@ -38,8 +34,6 @@ def can_use_musa_sampler( generators: dict[int, torch.Generator], logprobs_mode: LogprobsMode, ) -> bool: - if not sampler_fast_path_enabled(): - return False if not current_platform.is_musa() or not is_musa_tensor(logits): return False if generators: @@ -142,9 +136,7 @@ def sample_probs_seeded_multinomial( for row_idx in range(probs.shape[0]): generator = generators.get(row_idx) if generator is None: - sample = torch.multinomial( - probs[row_idx], num_samples=1, replacement=True - ) + sample = torch.multinomial(probs[row_idx], num_samples=1, replacement=True) else: sample = torch.multinomial( probs[row_idx], @@ -199,7 +191,7 @@ def _apply_top_k_top_p( if p is None and k is None: return logits - if current_platform.is_musa() and sampler_fast_path_enabled(): + if current_platform.is_musa(): if k is not None and logits.shape[0] >= 16: if p is None and logits.shape[1] >= 65536: max_top_k = int(k.to(torch.long).max().item()) @@ -264,7 +256,6 @@ def _topk_topp_sampler_init( if ( logprobs_mode not in ("processed_logits", "processed_logprobs") and current_platform.is_musa() - and sampler_fast_path_enabled() ): vllm_topk_topp_sampler.logger.info_once( "Using MUSA native ops for top-p/top-k/min-p sampling.", @@ -441,8 +432,6 @@ def can_use_worker_sampler( sampling_states: Any, idx_mapping_np: np.ndarray, ) -> bool: - if not sampler_fast_path_enabled(): - return False if not current_platform.is_musa() or not is_musa_tensor(logits): return False if logprobs_mode == "processed_logprobs": @@ -557,7 +546,9 @@ def _apply_worker_sampling_filters_for_seeded_multinomial( ) logits = _apply_top_k_top_p(logits, top_k, top_p) if use_min_p: - sampler.sampling_states.apply_min_p(logits, expanded_idx_mapping, idx_mapping_np) + sampler.sampling_states.apply_min_p( + logits, expanded_idx_mapping, idx_mapping_np + ) return logits @@ -571,11 +562,8 @@ def _worker_sample( expanded_local_pos: torch.Tensor, return_logprobs: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - if ( - logits.shape[0] == idx_mapping_np.shape[0] - and can_use_worker_seeded_multinomial( - logits, self.logprobs_mode, self.sampling_states, idx_mapping_np - ) + if logits.shape[0] == idx_mapping_np.shape[0] and can_use_worker_seeded_multinomial( + logits, self.logprobs_mode, self.sampling_states, idx_mapping_np ): vllm_topk_topp_sampler.logger.info_once( "Using MUSA seeded multinomial sampling for user-seeded requests.",