@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dc48ba0c75
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

Kernel-override Determinism [1/n] (#25603)

Signed-off-by: Bram Wasti <bwasti@meta.com>

Bram Wasti committed Sep 26, 2025 at 19:59 UTC dc48ba0c750e176f6314504cc0e8370a46ed01a8
8 files changed +890 -4
csrc/core/batch_invariant.hpp new
+16
@@ -0,0 +1,16 @@
1 +#pragma once
2 +#include <cstdlib>
3 +#include <string>
4 +#include <cctype>
5 +
6 +namespace vllm {
7 +
8 +// vllm_kernel_override_batch_invariant(); returns true
9 +// if env VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT=1
10 +inline bool vllm_kernel_override_batch_invariant() {
11 + std::string env_key = "VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT";
12 + const char* val = std::getenv(env_key.c_str());
13 + return (val && std::atoi(val) != 0) ? 1 : 0;
14 +}
15 +
16 +} // namespace vllm
csrc/layernorm_kernels.cu
+6 -2
@@ -1,6 +1,7 @@
1 #include "type_convert.cuh"
2 #include "dispatch_utils.h"
3 #include "cub_helpers.h"
4 +#include "core/batch_invariant.hpp"
5
6 #include <torch/cuda.h>
7 #include <c10/cuda/CUDAGuard.h>
@@ -413,7 +414,9 @@ void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]
414 wt_ptr % req_alignment_bytes == 0;
415 bool offsets_are_multiple_of_vector_width =
416 hidden_size % vector_width == 0 && input_stride % vector_width == 0;
416 - if (ptrs_are_aligned && offsets_are_multiple_of_vector_width) {
417 + bool batch_invariant_launch = vllm::vllm_kernel_override_batch_invariant();
418 + if (ptrs_are_aligned && offsets_are_multiple_of_vector_width &&
419 + !batch_invariant_launch) {
420 LAUNCH_FUSED_ADD_RMS_NORM(8);
421 } else {
422 LAUNCH_FUSED_ADD_RMS_NORM(0);
@@ -459,7 +462,8 @@ void poly_norm(torch::Tensor& out, // [..., hidden_size]
462 auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
463 auto out_ptr = reinterpret_cast<std::uintptr_t>(out.data_ptr());
464 bool ptrs_are_aligned = inp_ptr % 16 == 0 && out_ptr % 16 == 0;
462 - if (ptrs_are_aligned && hidden_size % 8 == 0) {
465 + bool batch_invariant_launch = vllm::vllm_kernel_override_batch_invariant();
466 + if (ptrs_are_aligned && hidden_size % 8 == 0 && !batch_invariant_launch) {
467 LAUNCH_FUSED_POLY_NORM(8);
468 } else {
469 LAUNCH_FUSED_POLY_NORM(0);
csrc/layernorm_quant_kernels.cu
+4 -1
@@ -9,6 +9,7 @@
9 #include "quantization/fp8/common.cuh"
10 #include "dispatch_utils.h"
11 #include "cub_helpers.h"
12 +#include "core/batch_invariant.hpp"
13
14 #include <torch/cuda.h>
15 #include <c10/cuda/CUDAGuard.h>
@@ -240,7 +241,9 @@ void fused_add_rms_norm_static_fp8_quant(
241 auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
242 bool ptrs_are_aligned =
243 inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;
243 - if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0) {
244 + bool batch_invariant_launch = vllm::vllm_kernel_override_batch_invariant();
245 + if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 &&
246 + !batch_invariant_launch) {
247 LAUNCH_FUSED_ADD_RMS_NORM(8);
248 } else {
249 LAUNCH_FUSED_ADD_RMS_NORM(0);
csrc/moe/topk_softmax_kernels.cu
+3 -1
@@ -21,6 +21,7 @@
21 #include <c10/cuda/CUDAGuard.h>
22 #include "../cuda_compat.h"
23 #include "../cub_helpers.h"
24 +#include "../core/batch_invariant.hpp"
25
26 #define MAX(a, b) ((a) > (b) ? (a) : (b))
27 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -405,7 +406,8 @@ void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, f
406 using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM>;
407 static constexpr int VPT = Constants::VPT;
408 static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
408 - const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
409 + const bool batch_invariant_launch = vllm::vllm_kernel_override_batch_invariant();
410 + const int num_warps = batch_invariant_launch ? 32 : (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
411 const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
412
413 dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
tests/v1/generation/test_batch_invariance.py new
+290
@@ -0,0 +1,290 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import contextlib
4 +import os
5 +import random
6 +import string
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm import LLM, SamplingParams
12 +
13 +
14 +def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
15 + # Lightweight random prompt generator to vary prompt lengths and content.
16 + vocab = [
17 + "alpha",
18 + "bravo",
19 + "charlie",
20 + "delta",
21 + "echo",
22 + "foxtrot",
23 + "golf",
24 + "hotel",
25 + "india",
26 + "juliet",
27 + "kilo",
28 + "lima",
29 + "mike",
30 + "november",
31 + "oscar",
32 + "papa",
33 + "quebec",
34 + "romeo",
35 + "sierra",
36 + "tango",
37 + "uniform",
38 + "victor",
39 + "whiskey",
40 + "xray",
41 + "yankee",
42 + "zulu",
43 + ]
44 + n = random.randint(min_words, max_words)
45 + words = random.choices(vocab, k=n)
46 +
47 + # Add some noise and punctuation variability
48 + if random.random() < 0.5:
49 + words[0] = words[0].capitalize()
50 + if random.random() < 0.2:
51 + words.append("".join(random.choices(string.ascii_lowercase, k=5)))
52 + punct = random.choice([".", "?", "!", "...", ""])
53 + return " ".join(words) + punct
54 +
55 +
56 +@pytest.mark.timeout(1000)
57 +def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():
58 + """
59 + Ensures that the same request (the 'needle' prompt) yields identical output
60 + whether run alone (bs=1) or mixed into a larger batch (e.g., bs=64),
61 + using the high-level v1 LLM() API only (no manual batching).
62 +
63 + Strategy:
64 + - Create two LLM engines with identical config except max_num_seqs: 1 vs N.
65 + - Compute a baseline output for the needle prompt with the bs=1 engine.
66 + - For many trials, generate a batch (size N) where the needle appears at a
67 + random position among random filler prompts using the bs=N engine.
68 + - Track how many trials match vs mismatch, and report totals at the end.
69 + The test fails if any mismatches occur, but we still dump pass/fail
70 + counts.
71 +
72 + Notes:
73 + - Use seeded stochastic sampling with a fixed seed to test determinism.
74 + - Outputs are intentionally longer and sampled at higher temperature/top_p
75 + to produce a more random-sounding phrase, yet remain deterministic by
76 + seed.
77 + - Keep max_tokens and max_model_len bounded for speed and memory use.
78 + """
79 + random.seed(12345)
80 +
81 + # Allow overrides from environment (useful for CI tuning)
82 + # "facebook/opt-125m" is too small, doesn't reliably test determinism
83 + model = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
84 + num_trials = int(os.getenv("VLLM_NEEDLE_TRIALS", "5"))
85 + batch_size = int(os.getenv("VLLM_NEEDLE_BATCH_SIZE", "64"))
86 + assert batch_size >= 2, "Batch size should be >= 2 to mix needle."
87 +
88 + # Keep GPU memory usage low to avoid startup allocation failures.
89 + gpu_mem_util = float(os.getenv("VLLM_GPU_MEMORY_UTILIZATION", "0.3"))
90 + max_model_len = int(os.getenv("VLLM_MAX_MODEL_LEN", "4096"))
91 + swap_space_gb = int(os.getenv("VLLM_SWAP_SPACE_GB", "4"))
92 +
93 + # Sampling parameters: longer outputs with a more random-sounding
94 + # continuation,but still deterministic due to fixed seed.
95 + temperature = float(os.getenv("VLLM_NEEDLE_TEMPERATURE", "0.0"))
96 + top_p = float(os.getenv("VLLM_NEEDLE_TOP_P", "0.95"))
97 + max_tokens = int(os.getenv("VLLM_NEEDLE_MAX_TOKENS", "128"))
98 +
99 + sampling = SamplingParams(
100 + temperature=temperature,
101 + top_p=top_p,
102 + max_tokens=max_tokens,
103 + seed=20240919,
104 + )
105 +
106 + needle_prompt = ("There once was a ")
107 +
108 + llm_bs1 = None
109 + llm_bsN = None
110 + try:
111 + # Engine with bs=1 behavior
112 + llm_bs1 = LLM_with_max_seqs(
113 + model=model,
114 + max_num_seqs=1,
115 + gpu_memory_utilization=gpu_mem_util,
116 + max_model_len=max_model_len,
117 + swap_space=swap_space_gb,
118 + )
119 +
120 + # Baseline generation for the needle prompt alone.
121 + baseline_out = llm_bs1.generate([needle_prompt], sampling)
122 + assert len(baseline_out) == 1
123 + assert len(baseline_out[0].outputs) >= 1
124 + baseline_text = baseline_out[0].outputs[0].text
125 +
126 + # Engine with larger batch limit (e.g., 64)
127 + llm_bsN = LLM_with_max_seqs(
128 + model=model,
129 + max_num_seqs=batch_size,
130 + gpu_memory_utilization=gpu_mem_util,
131 + max_model_len=max_model_len,
132 + swap_space=swap_space_gb,
133 + )
134 +
135 + mismatches = 0
136 +
137 + for trial in range(num_trials):
138 + # Create a batch of size `batch_size` and insert the needle at
139 + # a random index
140 + prompts: list[str] = []
141 + needle_pos = random.randint(0, batch_size - 1)
142 + for i in range(batch_size):
143 + if i == needle_pos:
144 + prompts.append(needle_prompt)
145 + else:
146 + prompts.append(_random_prompt())
147 +
148 + # Generate with the larger-batch engine
149 + outputs = llm_bsN.generate(prompts, sampling)
150 + # Find the needle output by position
151 + needle_output = outputs[needle_pos]
152 + assert needle_output.prompt == needle_prompt
153 + assert len(needle_output.outputs) >= 1
154 + text = needle_output.outputs[0].text
155 +
156 + if text != baseline_text:
157 + mismatches += 1
158 +
159 + passes = num_trials - mismatches
160 + # Dump how many passed vs failed
161 + print(f"[determinism] total={num_trials}, passed={passes}, "
162 + f"failed={mismatches}, batch_size={batch_size}")
163 +
164 + if mismatches > 0:
165 + pytest.fail(
166 + f"Nondeterministic outputs detected: {mismatches} failed out "
167 + f"of {num_trials} trials (batch_size={batch_size}).")
168 +
169 + finally:
170 + # Ensure engines are shutdown to free GPU/VRAM across test sessions
171 + if llm_bs1 is not None:
172 + with contextlib.suppress(Exception):
173 + llm_bs1.shutdown()
174 + if llm_bsN is not None:
175 + with contextlib.suppress(Exception):
176 + llm_bsN.shutdown()
177 +
178 +
179 +def _extract_step_logprobs(request_output):
180 + if getattr(request_output, "outputs", None):
181 + inner = request_output.outputs[0]
182 + if hasattr(inner, "logprobs") and inner.logprobs is not None:
183 + t = torch.tensor(
184 + [
185 + inner.logprobs[i][tid].logprob
186 + for i, tid in enumerate(inner.token_ids)
187 + ],
188 + dtype=torch.float32,
189 + )
190 + return t
191 +
192 + return None
193 +
194 +
195 +@pytest.mark.skipif(
196 + not torch.cuda.is_available(),
197 + reason="Requires CUDA to match production inference path.",
198 +)
199 +def test_logprobs_bitwise_batch_invariance_bs1_vs_bs2():
200 +
201 + #model_name = os.getenv("VLLM_TEST_MODEL", "facebook/opt-125m")
202 + model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
203 + tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
204 +
205 + # Force float32 to avoid precision-induced differences.
206 + llm = LLM(
207 + model=model_name,
208 + tensor_parallel_size=tp_size,
209 + enforce_eager=True, # helps reduce nondeterminism from some backends
210 + )
211 +
212 + prompts = [
213 + "The capital of France is",
214 + "The capital of Germany is",
215 + ]
216 +
217 + sp = SamplingParams(
218 + temperature=0.0,
219 + top_p=1.0,
220 + max_tokens=8,
221 + # Seed shouldn't matter at temperature=0, but keeping it stable anyway.
222 + seed=1234,
223 + logprobs=5,
224 + )
225 +
226 + # BS=1: run prompts individually and collect logprobs per step.
227 + bs1_logprobs_per_prompt = []
228 + for p in prompts:
229 + outs = llm.generate([p], sp, use_tqdm=False)
230 + assert len(outs) == 1
231 + step_logprobs = _extract_step_logprobs(outs[0])
232 + if step_logprobs is None:
233 + pytest.skip("Logits are not available on RequestOutput; "
234 + "enable logprobs return to run this test.")
235 + bs1_logprobs_per_prompt.append(step_logprobs)
236 +
237 + # BS=2: run prompts in a batch and collect logprobs per step for each
238 + # prompt.
239 + outs_batched = llm.generate(prompts, sp, use_tqdm=False)
240 + assert len(outs_batched) == len(prompts)
241 + bs2_logprobs_per_prompt = []
242 + for o in outs_batched:
243 + step_logprobs = _extract_step_logprobs(o)
244 + if step_logprobs is None:
245 + pytest.skip("Logits are not available on RequestOutput; "
246 + "enable logprobs return to run this test.")
247 + bs2_logprobs_per_prompt.append(step_logprobs)
248 +
249 + # Compare step-by-step logprobs for each prompt between BS=1 and BS=2 runs.
250 + for i, (logprobs_bs1, logprobs_bs2) in enumerate(
251 + zip(bs1_logprobs_per_prompt, bs2_logprobs_per_prompt)):
252 + assert len(logprobs_bs1) == len(logprobs_bs2), (
253 + f"Different number of generation steps for prompt index {i}: "
254 + f"{len(logprobs_bs1)} (BS=1) vs {len(logprobs_bs2)} (BS=2)")
255 + for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bs2)):
256 + assert a.shape == b.shape, (
257 + f"Logits shape mismatch at prompt {i}, step {t}: "
258 + f"{a.shape} vs {b.shape}")
259 + # Bitwise exact equality.
260 + assert torch.equal(
261 + a, b), (f"Bitwise logprobs mismatch at prompt {i}, step {t} "
262 + f"(dtype={a.dtype}, shape={a.shape}).")
263 +
264 +
265 +def LLM_with_max_seqs(
266 + model: str,
267 + max_num_seqs: int,
268 + gpu_memory_utilization: float,
269 + max_model_len: int,
270 + swap_space: int,
271 +) -> LLM:
272 + """
273 + Helper to construct an LLM with a specific max_num_seqs (batch-size limit)
274 + using the high-level v1 LLM API, while constraining memory usage.
275 + """
276 + return LLM(
277 + model=model,
278 + max_num_seqs=max_num_seqs,
279 + # Constrain GPU memory pool so test can run even on busy GPUs.
280 + gpu_memory_utilization=gpu_memory_utilization,
281 + # Keep KV cache footprint small while allowing longer outputs.
282 + max_model_len=max_model_len,
283 + # Allow some CPU offload if needed.
284 + swap_space=swap_space,
285 + # Keep things lean and CI-friendly.
286 + dtype="float16",
287 + # Single-GPU by default; override externally if desired.
288 + tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
289 + trust_remote_code=os.getenv("VLLM_TRUST_REMOTE_CODE", "0") == "1",
290 + )
vllm/model_executor/layers/batch_invariant.py new
+561
@@ -0,0 +1,561 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import contextlib
4 +import os
5 +from collections import namedtuple
6 +from collections.abc import Callable
7 +from typing import Any, Union
8 +
9 +import torch
10 +import triton
11 +import triton.language as tl
12 +
13 +
14 +def _matmul_launch_metadata(grid: Callable[..., Any], kernel: Any,
15 + args: dict[str, Any]) -> dict[str, Any]:
16 + ret = {}
17 + m, n, k = args["M"], args["N"], args["K"]
18 + ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}]"
19 + if "tiles_per_update" in args:
20 + ret["name"] = (f"{kernel.name} [M={m}, N={n}, K={k}, "
21 + f"tiles_per_update={args['tiles_per_update']:02}]")
22 + if "c_ptr" in args:
23 + bytes_per_elem = args["c_ptr"].element_size()
24 + else:
25 + bytes_per_elem = 1 if args["FP8_OUTPUT"] else 2
26 + ret[f"flops{bytes_per_elem * 8}"] = 2.0 * m * n * k
27 + ret["bytes"] = bytes_per_elem * (m * k + n * k + m * n)
28 + return ret
29 +
30 +
31 +@triton.jit
32 +def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS):
33 + group_id = tile_id // num_pid_in_group
34 + first_pid_m = group_id * GROUP_SIZE_M
35 + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
36 + pid_m = first_pid_m + (tile_id % group_size_m)
37 + pid_n = (tile_id % num_pid_in_group) // group_size_m
38 + return pid_m, pid_n
39 +
40 +
41 +@triton.jit(launch_metadata=_matmul_launch_metadata)
42 +def matmul_kernel_persistent(
43 + a_ptr,
44 + b_ptr,
45 + c_ptr, #
46 + bias_ptr,
47 + M,
48 + N,
49 + K, #
50 + stride_am,
51 + stride_ak,
52 + stride_bk,
53 + stride_bn,
54 + stride_cm,
55 + stride_cn,
56 + BLOCK_SIZE_M: tl.constexpr, #
57 + BLOCK_SIZE_N: tl.constexpr, #
58 + BLOCK_SIZE_K: tl.constexpr, #
59 + GROUP_SIZE_M: tl.constexpr, #
60 + NUM_SMS: tl.constexpr, #
61 + A_LARGE: tl.constexpr,
62 + B_LARGE: tl.constexpr,
63 + C_LARGE: tl.constexpr,
64 + HAS_BIAS: tl.constexpr,
65 +):
66 + start_pid = tl.program_id(axis=0)
67 + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
68 + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
69 + k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
70 + num_tiles = num_pid_m * num_pid_n
71 +
72 + tile_id_c = start_pid - NUM_SMS
73 +
74 + offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K)
75 + num_pid_in_group = GROUP_SIZE_M * num_pid_n
76 +
77 + for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True):
78 + pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m,
79 + GROUP_SIZE_M, NUM_SMS)
80 + start_m = pid_m * BLOCK_SIZE_M
81 + start_n = pid_n * BLOCK_SIZE_N
82 + offs_am = start_m + tl.arange(0, BLOCK_SIZE_M)
83 + offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N)
84 + if A_LARGE:
85 + offs_am = offs_am.to(tl.int64)
86 + if B_LARGE:
87 + offs_bn = offs_bn.to(tl.int64)
88 + offs_am = tl.where(offs_am < M, offs_am, 0)
89 + offs_bn = tl.where(offs_bn < N, offs_bn, 0)
90 + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M),
91 + BLOCK_SIZE_M)
92 + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N),
93 + BLOCK_SIZE_N)
94 +
95 + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
96 + for ki in range(k_tiles):
97 + if A_LARGE or B_LARGE:
98 + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(
99 + tl.int64)
100 + else:
101 + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
102 + a_ptrs = a_ptr + (offs_am[:, None] * stride_am +
103 + offs_k[None, :] * stride_ak)
104 + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk +
105 + offs_bn[None, :] * stride_bn)
106 +
107 + a = tl.load(a_ptrs,
108 + mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K,
109 + other=0.0)
110 + b = tl.load(b_ptrs,
111 + mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K,
112 + other=0.0)
113 + accumulator = tl.dot(a, b, accumulator)
114 +
115 + tile_id_c += NUM_SMS
116 + pid_m, pid_n = _compute_pid(tile_id_c, num_pid_in_group, num_pid_m,
117 + GROUP_SIZE_M, NUM_SMS)
118 + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
119 + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
120 + if C_LARGE:
121 + offs_cm = offs_cm.to(tl.int64)
122 + offs_cn = offs_cn.to(tl.int64)
123 + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[
124 + None, :]
125 + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
126 + if HAS_BIAS:
127 + bias_ptrs = bias_ptr + offs_cn
128 + bias = tl.load(bias_ptrs, mask=offs_cn < N,
129 + other=0.0).to(tl.float32)
130 + accumulator += bias
131 + if c_ptr.dtype.element_ty == tl.float8e4nv:
132 + c = accumulator.to(tl.float8e4nv)
133 + else:
134 + c = accumulator.to(tl.float16)
135 + tl.store(c_ptrs, c, mask=c_mask)
136 +
137 +
138 +def matmul_persistent(a: torch.Tensor,
139 + b: torch.Tensor,
140 + bias: Union[torch.Tensor, None] = None):
141 + # Check constraints.
142 + assert a.shape[1] == b.shape[0], "Incompatible dimensions"
143 + assert a.dtype == b.dtype, "Incompatible dtypes"
144 + assert bias is None or bias.dim() == 1, (
145 + "Currently assuming bias is 1D, let Horace know if you run into this")
146 + NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count
147 + M, K = a.shape
148 + K, N = b.shape
149 + dtype = a.dtype
150 + # Allocates output.
151 + c = torch.empty((M, N), device=a.device, dtype=dtype)
152 +
153 + # 1D launch kernel where each block gets its own program.
154 + def grid(META):
155 + return (min(
156 + NUM_SMS,
157 + triton.cdiv(M, META["BLOCK_SIZE_M"]) *
158 + triton.cdiv(N, META["BLOCK_SIZE_N"])), )
159 +
160 + configs = {
161 + torch.bfloat16: {
162 + "BLOCK_SIZE_M": 128,
163 + "BLOCK_SIZE_N": 128,
164 + "BLOCK_SIZE_K": 64,
165 + "GROUP_SIZE_M": 8,
166 + "num_stages": 3,
167 + "num_warps": 8,
168 + },
169 + torch.float16: {
170 + "BLOCK_SIZE_M": 128,
171 + "BLOCK_SIZE_N": 256,
172 + "BLOCK_SIZE_K": 64,
173 + "GROUP_SIZE_M": 8,
174 + "num_stages": 3,
175 + "num_warps": 8,
176 + },
177 + torch.float32: {
178 + "BLOCK_SIZE_M": 128,
179 + "BLOCK_SIZE_N": 128,
180 + "BLOCK_SIZE_K": 32,
181 + "GROUP_SIZE_M": 8,
182 + "num_stages": 3,
183 + "num_warps": 8,
184 + },
185 + }
186 + # print(a.device, b.device, c.device)
187 + matmul_kernel_persistent[grid](
188 + a,
189 + b,
190 + c, #
191 + bias,
192 + M,
193 + N,
194 + K, #
195 + a.stride(0),
196 + a.stride(1), #
197 + b.stride(0),
198 + b.stride(1), #
199 + c.stride(0),
200 + c.stride(1), #
201 + NUM_SMS=NUM_SMS, #
202 + A_LARGE=a.numel() > 2**31,
203 + B_LARGE=b.numel() > 2**31,
204 + C_LARGE=c.numel() > 2**31,
205 + HAS_BIAS=bias is not None,
206 + **configs[dtype],
207 + )
208 + return c
209 +
210 +
211 +@triton.jit
212 +def _log_softmax_kernel(
213 + input_ptr,
214 + output_ptr,
215 + input_row_stride,
216 + output_row_stride,
217 + n_cols,
218 + BLOCK_SIZE: tl.constexpr,
219 +):
220 + """
221 + Compute log_softmax along the last dimension of a 2D tensor.
222 + Each block handles one row of the input tensor.
223 + """
224 + # Get the row index for this block
225 + row_idx = tl.program_id(0).to(tl.int64)
226 +
227 + # Compute base pointers for input and output rows
228 + row_start_ptr = input_ptr + row_idx * input_row_stride
229 + output_row_start_ptr = output_ptr + row_idx * output_row_stride
230 +
231 + # Step 1: Find maximum value in the row for numerical stability
232 + max_val = -float("inf")
233 + for col_offset in range(0, n_cols, BLOCK_SIZE):
234 + col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
235 + mask = col_idx < n_cols
236 +
237 + # Load values
238 + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=-float("inf"))
239 +
240 + # Update maximum
241 + max_val = tl.max(tl.maximum(vals, max_val))
242 +
243 + # Step 2: Compute sum of exp(x - max_val)
244 + sum_exp = 0.0
245 + for col_offset in range(0, n_cols, BLOCK_SIZE):
246 + col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
247 + mask = col_idx < n_cols
248 +
249 + # Load values
250 + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0)
251 +
252 + # Compute exp(x - max_val) and accumulate
253 + exp_vals = tl.exp(vals - max_val)
254 + sum_exp += tl.sum(tl.where(mask, exp_vals, 0.0))
255 +
256 + # Compute log(sum_exp)
257 + log_sum_exp = tl.log(sum_exp)
258 +
259 + # Step 3: Compute final log_softmax values: x - max_val - log_sum_exp
260 + for col_offset in range(0, n_cols, BLOCK_SIZE):
261 + col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
262 + mask = col_idx < n_cols
263 +
264 + # Load values
265 + vals = tl.load(row_start_ptr + col_idx, mask=mask)
266 +
267 + # Compute log_softmax
268 + output = vals - max_val - log_sum_exp
269 +
270 + # Store results
271 + tl.store(output_row_start_ptr + col_idx, output, mask=mask)
272 +
273 +
274 +def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor:
275 + """
276 + Compute log_softmax using Triton kernel.
277 +
278 + Args:
279 + input: Input tensor
280 + dim: Dimension along which to compute log_softmax
281 + (only -1 or last dim supported)
282 + >> Stashed changes
283 + Returns:
284 + Tensor with log_softmax applied along the specified dimension
285 + """
286 + if dim != -1 and dim != input.ndim - 1:
287 + raise ValueError("This implementation only supports log_softmax along "
288 + "the last dimension")
289 +
290 + # Flatten all dimensions except the last one
291 + original_shape = input.shape
292 + input_2d = input.reshape(-1, input.shape[-1])
293 + input_2d = input_2d.contiguous()
294 +
295 + n_rows, n_cols = input_2d.shape
296 +
297 + # Allocate output tensor
298 + output = torch.empty_like(input_2d)
299 +
300 + # Choose block size based on the number of columns
301 + BLOCK_SIZE = 1024
302 +
303 + # Launch kernel with one block per row
304 + grid = (n_rows, )
305 + _log_softmax_kernel[grid](
306 + input_2d,
307 + output,
308 + input_2d.stride(0),
309 + output.stride(0),
310 + n_cols,
311 + BLOCK_SIZE=BLOCK_SIZE,
312 + )
313 + # Reshape output back to original shape
314 + return output.reshape(original_shape)
315 +
316 +
317 +@triton.jit
318 +def mean_kernel(
319 + input_ptr,
320 + output_ptr,
321 + input_stride0,
322 + input_stride1,
323 + input_stride2,
324 + output_stride0,
325 + output_stride1,
326 + M, # size before reduction dim
327 + N, # size of reduction dim
328 + K, # size after reduction dim
329 + BLOCK_SIZE: tl.constexpr,
330 +):
331 + """
332 + Kernel for computing mean along a single dimension.
333 + Input is viewed as (M, N, K) where N is the dimension being reduced.
334 + """
335 + # Program ID gives us which output element we're computing
336 + pid = tl.program_id(0)
337 +
338 + # Compute output indices
339 + m_idx = pid // K
340 + k_idx = pid % K
341 +
342 + # Bounds check
343 + if m_idx >= M or k_idx >= K:
344 + return
345 +
346 + # Accumulate sum across reduction dimension
347 + acc = 0.0
348 + for n_start in range(0, N, BLOCK_SIZE):
349 + n_offsets = n_start + tl.arange(0, BLOCK_SIZE)
350 + mask = n_offsets < N
351 +
352 + # Calculate input indices
353 + input_idx = m_idx * input_stride0 + n_offsets * input_stride1 \
354 + + k_idx * input_stride2
355 +
356 + # Load and accumulate
357 + vals = tl.load(input_ptr + input_idx, mask=mask, other=0.0)
358 + acc += tl.sum(vals)
359 +
360 + # Compute mean and store
361 + mean_val = acc / N
362 + output_idx = m_idx * output_stride0 + k_idx * output_stride1
363 + tl.store(output_ptr + output_idx, mean_val)
364 +
365 +
366 +def mean_dim(input: torch.Tensor,
367 + dim: int,
368 + keepdim: bool = False,
369 + dtype: Union[torch.dtype, None] = None) -> torch.Tensor:
370 + """
371 + Triton implementation of torch.mean with single dimension reduction.
372 +
373 + Args:
374 + input: Input tensor
375 + dim: Single dimension along which to compute mean
376 + keepdim: Whether to keep the reduced dimension
377 + dtype: Output dtype. If None, uses input dtype
378 + (or float32 for integer inputs)
379 +
380 + Returns:
381 + Tensor with mean values along specified dimension
382 + """
383 + # Validate inputs
384 + assert input.is_cuda, "Input must be a CUDA tensor"
385 + assert -input.ndim <= dim < input.ndim, (
386 + f"Invalid dimension {dim} for tensor with {input.ndim} dimensions")
387 +
388 + # Handle negative dim
389 + if dim < 0:
390 + dim = dim + input.ndim
391 +
392 + # Handle dtype
393 + if dtype is None:
394 + if input.dtype in [torch.int8, torch.int16, torch.int32, torch.int64]:
395 + dtype = torch.float32
396 + else:
397 + dtype = input.dtype
398 +
399 + # Convert input to appropriate dtype if needed
400 + if input.dtype != dtype:
401 + input = input.to(dtype)
402 +
403 + # Get input shape and strides
404 + shape = list(input.shape)
405 +
406 + # Calculate dimensions for kernel
407 + M = 1
408 + for i in range(dim):
409 + M *= shape[i]
410 +
411 + N = shape[dim]
412 +
413 + K = 1
414 + for i in range(dim + 1, len(shape)):
415 + K *= shape[i]
416 +
417 + # Reshape input to 3D view (M, N, K)
418 + input_3d = input.reshape(M, N, K)
419 +
420 + # Create output shape
421 + if keepdim:
422 + output_shape = shape.copy()
423 + output_shape[dim] = 1
424 + else:
425 + output_shape = shape[:dim] + shape[dim + 1:]
426 +
427 + # Create output tensor
428 + output = torch.empty(output_shape, dtype=dtype, device=input.device)
429 +
430 + # Reshape output for kernel
431 + if keepdim:
432 + output_2d = output.reshape(M, 1, K).squeeze(1)
433 + else:
434 + output_2d = output.reshape(M, K)
435 +
436 + # Launch kernel
437 + grid = (M * K, )
438 + BLOCK_SIZE = 1024
439 +
440 + mean_kernel[grid](
441 + input_3d,
442 + output_2d,
443 + input_3d.stride(0),
444 + input_3d.stride(1),
445 + input_3d.stride(2),
446 + output_2d.stride(0),
447 + output_2d.stride(1) if output_2d.ndim > 1 else 0,
448 + M,
449 + N,
450 + K,
451 + BLOCK_SIZE,
452 + )
453 +
454 + return output
455 +
456 +
457 +def mm_batch_invariant(a, b):
458 + return matmul_persistent(a, b)
459 +
460 +
461 +def addmm_batch_invariant(bias, a, b):
462 + return matmul_persistent(a, b, bias=bias)
463 +
464 +
465 +def _log_softmax_batch_invariant(input, dim, _half_to_float):
466 + assert not _half_to_float, "not implemented"
467 + return log_softmax(input, dim=dim)
468 +
469 +
470 +def mean_batch_invariant(input,
471 + dim,
472 + keepdim=False,
473 + dtype: Union[torch.dtype, None] = None):
474 + assert dtype is None or dtype == torch.float32, \
475 + f"unsupported dtype: {dtype}"
476 +
477 + result = input.to(torch.float32)
478 +
479 + # Sort dimensions to reduce from largest to smallest to handle shifting dims
480 + # during iterative reduction.
481 + sorted_dims = sorted([d % input.ndim for d in dim], reverse=True)
482 +
483 + # Iteratively apply a deterministic mean.
484 + for d in sorted_dims:
485 + result = mean_dim(result, dim=d, keepdim=True)
486 +
487 + if not keepdim:
488 + # Squeeze the reduced dimensions.
489 + for d in sorted_dims:
490 + result = result.squeeze(d)
491 +
492 + return result
493 +
494 +
495 +_batch_invariant_MODE = False
496 +_batch_invariant_LIB = None
497 +
498 +
499 +def is_batch_invariant_mode_enabled():
500 + return _batch_invariant_MODE
501 +
502 +
503 +def enable_batch_invariant_mode():
504 + global _batch_invariant_MODE, _batch_invariant_LIB
505 + if _batch_invariant_MODE:
506 + return
507 +
508 + _batch_invariant_MODE = True
509 + _batch_invariant_LIB = torch.library.Library("aten", "IMPL")
510 + _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")
511 + _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
512 + _batch_invariant_LIB.impl("aten::_log_softmax",
513 + _log_softmax_batch_invariant, "CUDA")
514 + _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA")
515 +
516 +
517 +def disable_batch_invariant_mode():
518 + global _batch_invariant_MODE, _batch_invariant_LIB
519 + if _batch_invariant_LIB is not None:
520 + _batch_invariant_LIB._destroy()
521 + _batch_invariant_MODE = False
522 + _batch_invariant_LIB = None
523 +
524 +
525 +@contextlib.contextmanager
526 +def set_batch_invariant_mode(enabled: bool = True):
527 + global _batch_invariant_MODE, _batch_invariant_LIB
528 + old_data = (_batch_invariant_MODE, _batch_invariant_LIB)
529 + if enabled:
530 + enable_batch_invariant_mode()
531 + else:
532 + disable_batch_invariant_mode()
533 + yield
534 + if _batch_invariant_LIB is not None:
535 + _batch_invariant_LIB._destroy()
536 + _batch_invariant_MODE, _batch_invariant_LIB = old_data
537 +
538 +
539 +AttentionBlockSize = namedtuple("AttentionBlockSize", ["block_m", "block_n"])
540 +
541 +
542 +def get_batch_invariant_attention_block_size() -> AttentionBlockSize:
543 + return AttentionBlockSize(block_m=16, block_n=16)
544 +
545 +
546 +def vllm_kernel_override_batch_invariant():
547 + env_key = "VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"
548 + is_overridden = False
549 + val = os.getenv(env_key, "0")
550 + try:
551 + is_overridden = int(val) != 0
552 + except ValueError:
553 + is_overridden = False
554 + return is_overridden
555 +
556 +
557 +def init_batch_invariance():
558 + # this will hit all the csrc overrides as well
559 + if vllm_kernel_override_batch_invariant():
560 + os.environ["VLLM_ATTENTION_BACKEND"] = "FLEX_ATTENTION"
561 + enable_batch_invariant_mode()
vllm/v1/attention/backends/flex_attention.py
+7
@@ -18,6 +18,8 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
18 is_quantized_kv_cache)
19 from vllm.config import VllmConfig
20 from vllm.logger import init_logger
21 +from vllm.model_executor.layers.batch_invariant import (
22 + vllm_kernel_override_batch_invariant)
23 from vllm.utils import cdiv, is_torch_equal_or_newer
24 from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
25 CommonAttentionMetadata)
@@ -839,6 +841,11 @@ def get_kernel_options(query, block_m, block_n,
841 kernel_options: dict[str, Union[int, bool]] = {
842 "FORCE_USE_FLEX_ATTENTION": True,
843 }
844 + if vllm_kernel_override_batch_invariant():
845 + kernel_options["BLOCK_M"] = 16
846 + kernel_options["BLOCK_N"] = 16
847 + kernel_options["IS_DIVISIBLE"] = False
848 + return kernel_options
849 if use_direct_build:
850 kernel_options["BLOCK_M"] = block_m
851 kernel_options["BLOCK_N"] = block_n
vllm/v1/worker/gpu_model_runner.py
+3
@@ -192,6 +192,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
192 from vllm.model_executor.models.utils import set_cpu_offload_max_bytes
193 set_cpu_offload_max_bytes(
194 int(self.cache_config.cpu_offload_gb * 1024**3))
195 + from vllm.model_executor.layers.batch_invariant import (
196 + init_batch_invariance)
197 + init_batch_invariance()
198
199 model_config = self.model_config
200 cache_config = self.cache_config