diff --git a/CMakeLists.txt b/CMakeLists.txt index d91f9a495e..c52eed869d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1739,7 +1739,8 @@ if(VLLM_CPP_HIP) src/vt/rocm/rocm_mla_ops.hip src/vt/rocm/rocm_mla_attn.hip src/vt/rocm/rocm_skinny_gemm.hip - src/vt/rocm/rocm_ops.hip) + src/vt/rocm/rocm_ops.hip + src/vt/rocm/rocm_combine_tokens.hip) if(VLLM_CPP_HIP_ARCHITECTURES) set_source_files_properties( src/vt/rocm/rocm_backend.hip @@ -1768,6 +1769,7 @@ if(VLLM_CPP_HIP) src/vt/rocm/rocm_mla_attn.hip src/vt/rocm/rocm_skinny_gemm.hip src/vt/rocm/rocm_ops.hip + src/vt/rocm/rocm_combine_tokens.hip PROPERTIES HIP_ARCHITECTURES "${VLLM_CPP_HIP_ARCHITECTURES}") endif() # Prefer the absolute path inside ${ROCM_PATH}/lib, fall back to the bare name, diff --git a/include/vt/rocm/combine_tokens.h b/include/vt/rocm/combine_tokens.h new file mode 100644 index 0000000000..aa65dc5b45 --- /dev/null +++ b/include/vt/rocm/combine_tokens.h @@ -0,0 +1,73 @@ +// Device combine/scatter kernels for async-scheduling overlap — ROCm/HIP port +// of include/vt/cuda/combine_tokens.h. Same contract: these replace the host +// scatter + its pre-sync (GPUModelRunner::sample_tokens_async's Synchronize +// before the host loop). The scatter writes last_sampled on the MAIN queue and +// the combine reads it on the MAIN queue, both stream-ordered relative to the +// forward, so no host round-trip of the sampled ids. +// +// Declarations only — the definitions live in src/vt/rocm/rocm_combine_tokens.hip. +// Signatures use plain pointers + vt::Queue so the header stays host-compilable. +#ifndef VT_ROCM_COMBINE_TOKENS_H_ +#define VT_ROCM_COMBINE_TOKENS_H_ + +#include + +#include "vt/backend.h" // vt::Queue + +namespace vt::rocm { + +// combine_sampled_and_draft_tokens (input_batch.py:304-406). Same FULL contract +// as the CUDA CombineKernel (vt/cuda/combine_tokens.h): num_logits comes from +// cu_num_logits (null == arange, i.e. ONE per row — NOT num_new_sampled_tokens; +// the two part at 0), and num_logits - num_new_sampled_tokens draft tokens are +// spliced from draft_tokens[req_state * draft_tokens_stride + b]. The current +// runner only ever reaches the subset draft_tokens == nullptr && +// cu_num_logits == nullptr && num_new_sampled_tokens == 1 (T0 non-spec), but +// the signatures must stay IDENTICAL across backends — the shared dispatcher +// (runner.cpp DispatchCombineSampledAndDraftTokens) forwards the same argument +// list to both, and the CUDA arm asserts the draft-bearing staging (A2-3) with +// a device trap rather than a silent skip. For each request row b, if the row +// is a decode row (seq_lens[b] > prefill_len[req_state]) splice the last +// sampled token into input_ids at the decode position (query_start_loc[b+1] - +// num_logits). Prefill/chunked-prefill rows (seq_len <= prefill_len) keep their +// prompt token. idx_mapping is the batch-row -> req_state indirection (the +// abort/finish reorder); pass nullptr for the identity mapping (our persistent +// batch is condensed dense, so batch row == req_state slot). Our runner builds +// logits_indices in prepare_inputs, so this kernel writes only the input_ids +// splice (the upstream kernel's logits_indices store is not needed here). +// Launched on the MAIN queue BEFORE the forward (outside any decode-graph +// capture — input prep always precedes the graph replay). +void LaunchCombineSampledAndDraftTokens( + Queue& queue, int32_t* input_ids, const int32_t* idx_mapping, + const int32_t* last_sampled_tokens, const int32_t* query_start_loc, + const int32_t* seq_lens, const int32_t* prefill_len, + const int32_t* draft_tokens, int draft_tokens_stride, + const int32_t* cu_num_logits, int num_reqs, int num_new_sampled_tokens); + +// post_update last_sampled scatter (input_batch.py:457-543 / states.py): record +// each row's freshly sampled id into last_sampled_tokens[req_state] on the MAIN +// queue, so the NEXT step's combine reads it without a sampled-id host +// round-trip. sampled_ids is the device-resident [num_reqs] argmax buffer the +// async sampler wrote (int64). idx_mapping is the batch-row -> req_state +// indirection (nullptr == identity). Replaces the runner's host scatter loop and +// its preceding Synchronize. +void LaunchScatterLastSampled(Queue& queue, int32_t* last_sampled_tokens, + const int64_t* sampled_ids, + const int32_t* idx_mapping, int num_reqs); + +// W4 (discrete GPU): replay InputBatch's STRUCTURAL edits to last_sampled_tokens +// onto the device mirror, in stream order. +// +// `ops` is a flat [4 * num_ops] int32 device array of (kind, a, b, value): +// kind 0 SEED: last_sampled[a] = value (add_request) +// kind 1 MOVE: last_sampled[a] = last_sampled[b] (condense) +// kind 2 SWAP: swap(last_sampled[a], last_sampled[b]) (swap_states) +// Applied STRICTLY IN ORDER by a single thread: the ops are not independent (a +// move can read a slot a previous move wrote), and there are at most a handful +// per step, so serial application is both correct and free. +void LaunchApplyLastSampledOps(Queue& queue, int32_t* last_sampled_tokens, + const int32_t* ops, int num_ops); + +} // namespace vt::rocm + +#endif // VT_ROCM_COMBINE_TOKENS_H_ diff --git a/src/vllm/v1/worker/gpu/runner.cpp b/src/vllm/v1/worker/gpu/runner.cpp index 2bd04cdf57..4fc6112707 100644 --- a/src/vllm/v1/worker/gpu/runner.cpp +++ b/src/vllm/v1/worker/gpu/runner.cpp @@ -48,6 +48,61 @@ #ifdef VLLM_CPP_CUDA #include "vt/cuda/combine_tokens.h" // W3 device combine/scatter (removes the sync) #endif +#ifdef VLLM_CPP_HIP +#include "vt/rocm/combine_tokens.h" // W3 device combine/scatter (ROCm port) +#endif + +// Device-agnostic dispatch for the combine/scatter/ops kernels. The CUDA and +// ROCm backends expose identical signatures in vt::cuda and vt::rocm; this +// dispatch compiles the right one in at build time. On a CPU-only build both +// backends are absent, so the entire block is compiled out — an empty anonymous +// namespace with three unused functions is a -Werror=unused-function break +// there (same treatment as AsyncDeviceMirrorEnvDefault below). No runtime +// device-type check: the build is single-backend, so the #if selects +// unconditionally. +#if defined(VLLM_CPP_CUDA) || defined(VLLM_CPP_HIP) +namespace { +void DispatchCombineSampledAndDraftTokens( + vt::Queue& q, int32_t* input_ids, const int32_t* idx_mapping, + const int32_t* last_sampled_tokens, const int32_t* query_start_loc, + const int32_t* seq_lens, const int32_t* prefill_len, + const int32_t* draft_tokens, int draft_tokens_stride, + const int32_t* cu_num_logits, int num_reqs, int num_new_sampled_tokens) { +#if defined(VLLM_CPP_CUDA) + vt::cuda::LaunchCombineSampledAndDraftTokens( + q, input_ids, idx_mapping, last_sampled_tokens, query_start_loc, + seq_lens, prefill_len, draft_tokens, draft_tokens_stride, cu_num_logits, + num_reqs, num_new_sampled_tokens); +#elif defined(VLLM_CPP_HIP) + vt::rocm::LaunchCombineSampledAndDraftTokens( + q, input_ids, idx_mapping, last_sampled_tokens, query_start_loc, + seq_lens, prefill_len, draft_tokens, draft_tokens_stride, cu_num_logits, + num_reqs, num_new_sampled_tokens); +#endif +} + +void DispatchScatterLastSampled(vt::Queue& q, int32_t* last_sampled_tokens, + const int64_t* sampled_ids, + const int32_t* idx_mapping, int num_reqs) { +#if defined(VLLM_CPP_CUDA) + vt::cuda::LaunchScatterLastSampled(q, last_sampled_tokens, sampled_ids, + idx_mapping, num_reqs); +#elif defined(VLLM_CPP_HIP) + vt::rocm::LaunchScatterLastSampled(q, last_sampled_tokens, sampled_ids, + idx_mapping, num_reqs); +#endif +} + +void DispatchApplyLastSampledOps(vt::Queue& q, int32_t* last_sampled_tokens, + const int32_t* ops, int num_ops) { +#if defined(VLLM_CPP_CUDA) + vt::cuda::LaunchApplyLastSampledOps(q, last_sampled_tokens, ops, num_ops); +#elif defined(VLLM_CPP_HIP) + vt::rocm::LaunchApplyLastSampledOps(q, last_sampled_tokens, ops, num_ops); +#endif +} +} // namespace +#endif // VLLM_CPP_CUDA || VLLM_CPP_HIP namespace vllm::v1 { @@ -102,12 +157,14 @@ static bool AsyncRunnerEnvDefault() { // steps. Whether that read is VALID is a backend CAPABILITY, not a device name: // ask the backend (vt::Backend::SupportsAsyncSampledTokenReadback, backend.h), // which answers true for CPU (host and device memory are one allocation) and -// CUDA (the sampled id is device-mirrored, async_device_mirror()), and false for -// a DISCRETE non-CUDA GPU (e.g. ROCm gfx1201) whose sample_tokens_async leg -// host-dereferences a device Alloc — the root cause of the "!" tokens on the lab -// R9700 (2026-08-07). An absent backend (device not built into this binary) -// yields nullptr and therefore false, which also subsumes the old -// #ifdef VLLM_CPP_CUDA guard. Keeping the question on the backend is what stops +// CUDA and ROCm (the sampled id is device-mirrored, async_device_mirror() — the +// W3/W4 combine/scatter kernels ported to both backends), and false for a +// DISCRETE GPU whose backend has not ported those kernels, whose +// sample_tokens_async leg host-dereferences a device Alloc — the root cause of +// the "!" tokens on the lab R9700 (2026-08-07). An absent backend (device not +// built into this binary) yields nullptr and therefore false, which also +// subsumes the old #if defined(VLLM_CPP_CUDA) || defined(VLLM_CPP_HIP) guard. +// Keeping the question on the backend is what stops // this device-agnostic shared layer from naming a device (check-device-leakage). static bool QueueSupportsAsyncInputCombine(const vt::Queue& queue) { const vt::Backend* backend = vt::TryGetBackend(queue.device.type); @@ -2611,7 +2668,7 @@ std::optional GPUModelRunner::execute_model( "verify step splices the committed token over the LAST DRAFT " "SLOT and scatters no drafts — SPEC-DFLASH2 reason A, which " "costs acceptance and raises nothing)"); -#ifdef VLLM_CPP_CUDA +#if defined(VLLM_CPP_CUDA) || defined(VLLM_CPP_HIP) // The host-side source both device arms start from. It is NOT true that they // therefore cannot be edited apart, and the round-3 review was right to say // so: only the UMA arm passes this pointer through. The mirror arm passes @@ -2678,7 +2735,7 @@ std::optional GPUModelRunner::execute_model( // which is reason A on any verify step — see the block above the #ifdef. stage_upload(*dev, dev->cu_num_logits, cu_num_logits_host, static_cast(num_reqs) + 1); - vt::cuda::LaunchCombineSampledAndDraftTokens( + DispatchCombineSampledAndDraftTokens( queue_, dev->input_ids, /*idx_mapping=*/nullptr, dev->last_sampled, dev->query_start_loc, dev->seq_lens, dev->prefill_len, dev->draft_tokens, dev->draft_stride, @@ -2715,7 +2772,7 @@ std::optional GPUModelRunner::execute_model( // be null here, which the kernel reads as arange; this is GB10's // production default, so that null was reason A on the integrated // production path from the first verify step A2-5 would admit. - vt::cuda::LaunchCombineSampledAndDraftTokens( + DispatchCombineSampledAndDraftTokens( queue_, step.input_token_ids.data(), /*idx_mapping=*/nullptr, input_batch_.last_sampled_tokens.data(), step.query_start_loc.data(), step.seq_lens.data(), input_batch_.prefill_len.data(), @@ -5288,7 +5345,7 @@ AsyncOutputPool& GPUModelRunner::get_or_create_async_output_pool() { // Distinct from VT_ASYNC_RUNNER, which would also turn off async scheduling // itself; keeping them separate is what makes an honest A/B of W4 alone possible — // same binary, same scheduler, one mechanism. -#ifdef VLLM_CPP_CUDA +#if defined(VLLM_CPP_CUDA) || defined(VLLM_CPP_HIP) // Guarded with its only use below: on a CPU build the mirror cannot exist, and // an unused static function is a -Werror=unused-function break there. DEFAULT ON: // on unless VT_ASYNC_DEVICE_MIRROR is explicitly "0" (the rollback), mirroring the @@ -5302,7 +5359,7 @@ static bool AsyncDeviceMirrorEnvDefault() { bool GPUModelRunner::async_device_mirror() const { if (async_device_mirror_cached_ >= 0) return async_device_mirror_cached_ != 0; bool on = false; -#ifdef VLLM_CPP_CUDA +#if defined(VLLM_CPP_CUDA) || defined(VLLM_CPP_HIP) // Engage on any real CUDA GPU, integrated OR discrete — NOT the CPU backend. // - DISCRETE (separate memory, !UnifiedMemory): the mirror is REQUIRED, because // the host fallback would main-stream Synchronize to read the sampled ids. @@ -5336,7 +5393,7 @@ bool GPUModelRunner::async_device_mirror() const { bool GPUModelRunner::async_executor() const { if (async_executor_cached_ >= 0) return async_executor_cached_ != 0; bool on = false; -#ifdef VLLM_CPP_CUDA +#if defined(VLLM_CPP_CUDA) || defined(VLLM_CPP_HIP) const char* value = std::getenv("VT_ASYNC_EXECUTOR"); on = value != nullptr && value[0] == '1' && value[1] == '\0' && async_device_mirror(); @@ -5419,7 +5476,7 @@ void GPUModelRunner::stage_upload(AsyncDeviceInputs& dev, int32_t* dst, } void GPUModelRunner::replay_last_sampled_ops(AsyncDeviceInputs& dev) { -#ifdef VLLM_CPP_CUDA +#if defined(VLLM_CPP_CUDA) || defined(VLLM_CPP_HIP) std::vector& ops = input_batch_.last_sampled_ops; if (ops.empty()) return; // Flatten to (kind, a, b, value) quads. The log is bounded by the number of @@ -5444,7 +5501,7 @@ void GPUModelRunner::replay_last_sampled_ops(AsyncDeviceInputs& dev) { VT_CHECK(static_cast(flat.size()) <= cap_ops, "async device mirror: structural-op chunk exceeds its buffer"); stage_upload(dev, dev.ops, flat.data(), static_cast(flat.size())); - vt::cuda::LaunchApplyLastSampledOps(queue_, dev.last_sampled, dev.ops, + DispatchApplyLastSampledOps(queue_, dev.last_sampled, dev.ops, static_cast(chunk)); done += chunk; } @@ -5579,7 +5636,7 @@ std::unique_ptr GPUModelRunner::sample_tokens_async( // a fourth one cannot forget to. The branches keep their write-backs (the // mirror and the device array are what the NEXT step's combine reads); what // they no longer own is the propose's view of them. -#ifdef VLLM_CPP_CUDA +#if defined(VLLM_CPP_CUDA) || defined(VLLM_CPP_HIP) // W4 device-resident scatter. Preferred whenever the mirror is engaged // (async_device_mirror(): CUDA + VT_ASYNC_DEVICE_MIRROR, INTEGRATED OR DISCRETE): // write each row's sampled id into the DEVICE mirror (dinp->last_sampled) on the @@ -5593,7 +5650,7 @@ std::unique_ptr GPUModelRunner::sample_tokens_async( // async output's own copy, as upstream does). Runs OUTSIDE any CUDA-graph capture. if (AsyncDeviceInputs* dinp = get_or_create_async_device_inputs(); dinp != nullptr) { - vt::cuda::LaunchScatterLastSampled(queue_, dinp->last_sampled, + DispatchScatterLastSampled(queue_, dinp->last_sampled, static_cast(dev_ids), /*idx_mapping=*/nullptr, num_reqs); for (int i = 0; i < num_reqs; ++i) { @@ -5617,7 +5674,7 @@ std::unique_ptr GPUModelRunner::sample_tokens_async( // this is the array condense reorders, its scatter pins the drain to // execute_model's top (the mirror path lifts that). is_integrated_gpu() // decouples a future discrete GPU (false -> host bookkeeping below). - vt::cuda::LaunchScatterLastSampled( + DispatchScatterLastSampled( queue_, input_batch_.last_sampled_tokens.data(), static_cast(dev_ids), /*idx_mapping=*/nullptr, num_reqs); for (int i = 0; i < num_reqs; ++i) { diff --git a/src/vt/rocm/rocm_backend.hip b/src/vt/rocm/rocm_backend.hip index 42eea3d3b6..8f83ab6574 100644 --- a/src/vt/rocm/rocm_backend.hip +++ b/src/vt/rocm/rocm_backend.hip @@ -335,6 +335,14 @@ class RocmBackend final : public Backend { // an illegal op during capture a loud failure on THIS thread rather than a // process-wide mode change. bool SupportsGraphCapture() const override { return true; } + // The sampled token id is device-mirrored (async_device_mirror(), the ROCm + // port of the W3/W4 combine/scatter kernels in rocm_combine_tokens.hip), so + // the between-steps host readback the depth-2 async input-combine needs is + // valid. Mirrors the CUDA backend's override (cuda_backend.cu:229-230); + // without this the ported kernels are dead code — runner_supports_async() + // resolves false and async scheduling stays off (AGENTS.md "Nothing lands + // dead"). + bool SupportsAsyncSampledTokenReadback() const override { return true; } void BeginCapture(Queue& q) override { Check(hipStreamBeginCapture(AsStream(q), hipStreamCaptureModeThreadLocal), "hipStreamBeginCapture"); diff --git a/src/vt/rocm/rocm_combine_tokens.hip b/src/vt/rocm/rocm_combine_tokens.hip new file mode 100644 index 0000000000..04e132658e --- /dev/null +++ b/src/vt/rocm/rocm_combine_tokens.hip @@ -0,0 +1,157 @@ +// ROCm/HIP port of src/vt/cuda/cuda_combine_tokens.cu — device combine/scatter +// kernels for async-scheduling overlap. Same logic, HIP API. These replace the +// host scatter + its pre-sync so the sampled ids never round-trip the host. +// +// The kernels are main-stream-ordered relative to the forward; on a discrete +// ROCm GPU the pointers are device allocations (the AsyncDeviceInputs mirror). +#include + +#include +#include +#include + +#include "vt/rocm/combine_tokens.h" + +namespace vt::rocm { +namespace { + +constexpr int kBlock = 256; + +void Check(hipError_t err, const char* what) { + if (err != hipSuccess) { + throw std::runtime_error(std::string("vt rocm combine_tokens: ") + what + + ": " + hipGetErrorString(err)); + } +} + +hipStream_t AsStream(const Queue& q) { + return static_cast(q.handle); +} + +// _combine_sampled_and_draft_tokens_kernel (input_batch.py:304-360), input_ids +// splice only (our logits_indices come from prepare_inputs). One thread per +// request row. Line-for-line the CUDA CombineKernel (cuda_combine_tokens.cu), +// full contract: num_logits derives from cu_num_logits (null == arange == ONE +// per row — not num_new_sampled_tokens; the two part at 0), and the +// num_logits - num_new_sampled_tokens draft tokens splice from the draft +// buffer. The trap replaces the host's two VT_CHECKs (stride narrower than the +// row count; null draft buffer while drafts are due) exactly as the CUDA arm's +// __trap() does — __builtin_trap() because HIP exposes no __trap() alias in +// every header set this file compiles under. +__global__ void CombineKernel(int32_t* input_ids, const int32_t* idx_mapping, + const int32_t* last_sampled_tokens, + const int32_t* query_start_loc, + const int32_t* seq_lens, + const int32_t* prefill_len, + const int32_t* draft_tokens, + int draft_tokens_stride, + const int32_t* cu_num_logits, int num_reqs, + int num_new_sampled_tokens) { + const int batch_idx = blockIdx.x * blockDim.x + threadIdx.x; + if (batch_idx >= num_reqs) return; + + const int req_state_idx = + idx_mapping != nullptr ? idx_mapping[batch_idx] : batch_idx; + + const int num_logits = + cu_num_logits != nullptr + ? cu_num_logits[batch_idx + 1] - cu_num_logits[batch_idx] + : 1; + const int num_draft_tokens = num_logits - num_new_sampled_tokens; + + const int query_end = query_start_loc[batch_idx + 1]; + const int logits_start = query_end - num_logits; + + const int seq_len = seq_lens[batch_idx]; + const int pf = prefill_len[req_state_idx]; + if (seq_len <= pf) return; + + const int first_logit_seq_pos = seq_len - num_logits; + if (num_new_sampled_tokens > 0 && first_logit_seq_pos >= pf) { + input_ids[logits_start] = last_sampled_tokens[req_state_idx]; + } + + if (num_draft_tokens > 0) { + if (draft_tokens == nullptr || draft_tokens_stride < num_draft_tokens) { + __builtin_trap(); + } + for (int b = 0; b < num_draft_tokens; ++b) { + input_ids[query_end - num_draft_tokens + b] = + draft_tokens[req_state_idx * draft_tokens_stride + b]; + } + } +} + +// post_update last_sampled scatter (input_batch.py:457-543 / states.py): one +// thread per request row writes the freshly sampled id into last_sampled_tokens. +__global__ void ScatterLastSampledKernel(int32_t* last_sampled_tokens, + const int64_t* sampled_ids, + const int32_t* idx_mapping, + int num_reqs) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_reqs) return; + const int req_state_idx = idx_mapping != nullptr ? idx_mapping[i] : i; + last_sampled_tokens[req_state_idx] = static_cast(sampled_ids[i]); +} + +// W4 structural replay (see combine_tokens.h). ONE thread, strictly in order: +// the ops are dependent (a condense move can read a slot an earlier move wrote) +// and there are a handful per step at most. +__global__ void ApplyLastSampledOpsKernel(int32_t* last_sampled_tokens, + const int32_t* ops, int num_ops) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + for (int i = 0; i < num_ops; ++i) { + const int32_t kind = ops[4 * i + 0]; + const int32_t a = ops[4 * i + 1]; + const int32_t b = ops[4 * i + 2]; + const int32_t value = ops[4 * i + 3]; + if (kind == 0) { + last_sampled_tokens[a] = value; + } else if (kind == 1) { + last_sampled_tokens[a] = last_sampled_tokens[b]; + } else if (kind == 2) { + const int32_t tmp = last_sampled_tokens[a]; + last_sampled_tokens[a] = last_sampled_tokens[b]; + last_sampled_tokens[b] = tmp; + } + } +} + +} // namespace + +void LaunchCombineSampledAndDraftTokens( + Queue& queue, int32_t* input_ids, const int32_t* idx_mapping, + const int32_t* last_sampled_tokens, const int32_t* query_start_loc, + const int32_t* seq_lens, const int32_t* prefill_len, + const int32_t* draft_tokens, int draft_tokens_stride, + const int32_t* cu_num_logits, int num_reqs, int num_new_sampled_tokens) { + if (num_reqs <= 0) return; + const int grid = (num_reqs + kBlock - 1) / kBlock; + hipLaunchKernelGGL(CombineKernel, dim3(grid), dim3(kBlock), 0, AsStream(queue), + input_ids, idx_mapping, last_sampled_tokens, + query_start_loc, seq_lens, prefill_len, draft_tokens, + draft_tokens_stride, cu_num_logits, num_reqs, + num_new_sampled_tokens); + Check(hipGetLastError(), "CombineKernel launch"); +} + +void LaunchScatterLastSampled(Queue& queue, int32_t* last_sampled_tokens, + const int64_t* sampled_ids, + const int32_t* idx_mapping, int num_reqs) { + if (num_reqs <= 0) return; + const int grid = (num_reqs + kBlock - 1) / kBlock; + hipLaunchKernelGGL(ScatterLastSampledKernel, dim3(grid), dim3(kBlock), 0, + AsStream(queue), last_sampled_tokens, sampled_ids, + idx_mapping, num_reqs); + Check(hipGetLastError(), "ScatterLastSampledKernel launch"); +} + +void LaunchApplyLastSampledOps(Queue& queue, int32_t* last_sampled_tokens, + const int32_t* ops, int num_ops) { + if (num_ops <= 0) return; + hipLaunchKernelGGL(ApplyLastSampledOpsKernel, dim3(1), dim3(1), 0, + AsStream(queue), last_sampled_tokens, ops, num_ops); + Check(hipGetLastError(), "ApplyLastSampledOpsKernel launch"); +} + +} // namespace vt::rocm