@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 18f8aa0465

[Refactor] Remove dead kernel code (#56845)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Wentao Ye committed Sep 15, 2026 at 16:17 UTC 18f8aa04656de0a37eb5c7e9efc8ebb595ac73cf
8 files changed -1895
vllm/model_executor/layers/mamba/ops/replayssm_config.py
-15
@@ -8,7 +8,6 @@ auto-detected (Blackwell vs not) so call sites need not thread it through.
8 """
9
10 import functools
11 -from contextlib import contextmanager
11
12 from vllm.platforms import current_platform
13 from vllm.triton_utils import triton
@@ -26,20 +25,6 @@ def _is_blackwell() -> bool:
25 _overrides: dict[str, tuple] = {}
26
27
29 -@contextmanager
30 -def override_replayssm_config(kernel: str, config: tuple):
31 - """Pin ``kernel``'s launch config for the duration of the context."""
32 - prev = _overrides.get(kernel)
33 - _overrides[kernel] = config
34 - try:
35 - yield
36 - finally:
37 - if prev is None:
38 - _overrides.pop(kernel, None)
39 - else:
40 - _overrides[kernel] = prev
41 -
42 -
28 def _dstate_tile(dstate: int, tile: int) -> int:
29 return max(16, min(tile, triton.next_power_of_2(dstate)))
30
vllm/models/glm5next/amd/ops/kpool_compress.py
-44
@@ -65,50 +65,6 @@ def _hadamard128(x):
65 return x * 0.08838834764831845 # 1/sqrt(128)
66
67
68 -# Map pool ids to physical cache slots.
69 -
70 -
71 -def compute_pooled_write_locs(
72 - page_table_64: torch.Tensor,
73 - pool_ids: torch.Tensor,
74 - pool_size: int,
75 -) -> torch.Tensor:
76 - """Map logical pooled-K ids to physical flat cache slots.
77 -
78 - ``pool_size`` consecutive tokens share one pool slot that lives at the
79 - *first* token page of each page-group. ``page_table_64`` maps token pages
80 - to physical block ids; we gather the block id of each pool's page-group
81 - and add the in-block pool offset.
82 - """
83 - assert page_table_64.ndim == 1
84 - pool_ids = pool_ids.to(torch.int64)
85 - block_size = 64
86 - pool_page_group = torch.div(pool_ids, block_size, rounding_mode="floor")
87 - token_page_row = pool_page_group * pool_size
88 - packed_page = page_table_64.index_select(0, token_page_row.to(torch.int64))
89 - return packed_page.to(torch.int64) * block_size + torch.remainder(
90 - pool_ids, block_size
91 - )
92 -
93 -
94 -def build_pooled_page_table(
95 - page_table: torch.Tensor,
96 - pool_size: int,
97 -) -> torch.Tensor:
98 - """Build a pool-granular page table by taking every ``pool_size``-th
99 - token-page column (one pool maps to ``pool_size`` token pages).
100 -
101 - Uses gather (not strided slicing) so the result is always a fresh
102 - row-major tensor — some downstream kernels require stride(-1) == 1.
103 - """
104 - block_size = page_table.shape[-1]
105 - assert block_size % pool_size == 0, (
106 - f"pool_size ({pool_size}) must divide page columns ({block_size})"
107 - )
108 - idx = torch.arange(0, block_size, pool_size, device=page_table.device)
109 - return page_table[..., idx].contiguous()
110 -
111 -
68 # Fused pool compression and cache write.
69
70
vllm/models/inkling/amd/ops/lamport.py deleted
-766
@@ -1,766 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm.
4 -
5 -The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``.
6 -
7 -Liveness
8 ---------
9 -Large grids are deliberately split at the two communication dependencies:
10 -
11 - 1. ``_publish_input_kernel`` only publishes rank partials.
12 - 2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard.
13 - 3. ``_sconv_publish_kernel`` only computes and publishes the local result.
14 - 4. ``_gather_norm_kernel`` waits, gathers, and normalizes.
15 -
16 -Without PDL, CUDA stream order completes a producer before its consumer. With
17 -PDL, every producer CTA posts all peer stores before triggering its dependent,
18 -and the consumer executes ``gdc_wait`` before polling. A consumer therefore
19 -waits only for stores from a producer that is already running (or complete) on
20 -another GPU. It never waits for another block in its own grid. Consequently
21 -spinning consumers cannot occupy resources needed by any producer, and the
22 -wait-for graph has no cycle. The proof is independent of grid size, block
23 -dispatch order, and occupancy.
24 -
25 -For one token, the first three phases use eight independent channel slices to
26 -expose enough CTA parallelism for decode latency. Each slice has exclusive
27 -ownership of its cache and Lamport columns; the gather/RMSNorm phase retains
28 -one CTA per token so no cross-CTA reduction or completion counter is needed.
29 -
30 -Immediate buffer reuse (including replay of a captured CUDA graph) is also
31 -safe. Rank R cannot republish input for call n+1 until its gather for n has
32 -finished; that gather waited for owner O's output, which O publishes only
33 -after consuming R's input for n. Likewise, R cannot republish output for n+1
34 -until its reduction for n+1 has observed destination D's input; D publishes
35 -that input only after its gather consumed R's output for n. Thus every prior
36 -read happens-before a same-slot rewrite. Three generations reduce incidental
37 -coupling for ordinary launches, but correctness does not depend on rotation.
38 -
39 -The payload itself is the Lamport flag. A 32-bit store publishes two bf16s
40 -atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real
41 -negative zeroes are changed to positive zero before publication. Consumers
42 -use volatile 32-bit loads and restore the sentinel after consuming a slot.
43 -"""
44 -
45 -from __future__ import annotations
46 -
47 -import os
48 -
49 -import torch
50 -
51 -from vllm.distributed import get_tp_group
52 -from vllm.distributed.parallel_state import in_the_same_node_as
53 -from vllm.logger import init_logger
54 -from vllm.triton_utils import tl, triton
55 -
56 -logger = init_logger(__name__)
57 -
58 -_MAX_TOKENS = 16384
59 -_EMPTY_PAIR = tl.constexpr(0x80008000)
60 -
61 -
62 -@triton.jit
63 -def _pack_bf16_pairs(values):
64 - """Pack bf16 values into atomic u32 pairs and reserve negative zero."""
65 - lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2]))
66 - lo = lo.to(tl.uint16, bitcast=True)
67 - hi = hi.to(tl.uint16, bitcast=True)
68 - lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32)
69 - hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32)
70 - return lo | (hi << 16)
71 -
72 -
73 -@triton.jit
74 -def _unpack_bf16_pairs(values):
75 - lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
76 - hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
77 - return tl.interleave(lo, hi)
78 -
79 -
80 -@triton.jit
81 -def _wait_pairs(ptr, offsets, mask):
82 - values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True)
83 - while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0:
84 - values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True)
85 - return values
86 -
87 -
88 -@triton.jit
89 -def _publish_input_kernel(
90 - stage_ptr,
91 - peer_ptrs,
92 - peer_offset_u32,
93 - stride_stage_t,
94 - C: tl.constexpr,
95 - CS: tl.constexpr,
96 - CS_P2: tl.constexpr,
97 - SPLITS: tl.constexpr,
98 - RANK: tl.constexpr,
99 - WORLD: tl.constexpr,
100 - USE_PDL: tl.constexpr,
101 - launch_pdl: tl.constexpr,
102 -):
103 - """Publish this rank's full partial row into every shard owner's slots."""
104 - token = tl.program_id(0).to(tl.int64)
105 - split = tl.program_id(1).to(tl.int64)
106 - CSS: tl.constexpr = CS // SPLITS
107 - pair = tl.arange(0, CS_P2 // 2)
108 - pair_mask = pair < CSS // 2
109 - elem = tl.arange(0, CS_P2)
110 - elem_mask = elem < CSS
111 - ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64))
112 - # Address generation is independent of the preceding stream kernel. The
113 - # acquire remains before stage/peer loads, which may consume its writes.
114 - if USE_PDL:
115 - tl.extra.cuda.gdc_wait()
116 -
117 - for owner in tl.static_range(WORLD):
118 - values = tl.load(
119 - stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem,
120 - mask=elem_mask,
121 - other=0.0,
122 - )
123 - packed = _pack_bf16_pairs(values)
124 - base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32))
125 - dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair
126 - tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask)
127 - if USE_PDL:
128 - tl.extra.cuda.gdc_launch_dependents()
129 -
130 -
131 -@triton.jit
132 -def _reduce_insert_kernel(
133 - input_peer_ptrs,
134 - input_peer_offset_u32,
135 - cache_ptr,
136 - slot_ptr,
137 - stride_cache_block,
138 - stride_cache_head,
139 - stride_cache_token,
140 - stride_cache_dim,
141 - block_size,
142 - C: tl.constexpr,
143 - CS: tl.constexpr,
144 - CS_P2: tl.constexpr,
145 - SPLITS: tl.constexpr,
146 - HEAD_SIZE: tl.constexpr,
147 - CACHE_OFFSET: tl.constexpr,
148 - RANK: tl.constexpr,
149 - WORLD: tl.constexpr,
150 - USE_PDL: tl.constexpr,
151 - launch_pdl: tl.constexpr,
152 -):
153 - """Consume all partials for this rank and insert the reduced cache row."""
154 - token = tl.program_id(0).to(tl.int64)
155 - split = tl.program_id(1).to(tl.int64)
156 - CSS: tl.constexpr = CS // SPLITS
157 - source = tl.arange(0, WORLD)
158 - pair = tl.arange(0, CS_P2 // 2)
159 - pair_mask = pair < CSS // 2
160 - offsets = (
161 - (token * WORLD + source)[:, None] * (CS // 2)
162 - + split * (CSS // 2)
163 - + pair[None, :]
164 - )
165 - mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :]
166 - input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64))
167 - input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32))
168 - input_u32 += input_peer_offset_u32
169 - # Slot metadata and the cache destination do not depend on input publish.
170 - slot = tl.load(slot_ptr + token)
171 - valid = slot >= 0
172 - channel = tl.arange(0, CS_P2)
173 - channel_mask = channel < CSS
174 - global_channel = split * CSS + channel
175 - head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1)
176 - dim = CACHE_OFFSET + global_channel % HEAD_SIZE
177 - safe_slot = tl.maximum(slot, 0).to(tl.int64)
178 - dst = (
179 - cache_ptr
180 - + (safe_slot // block_size) * stride_cache_block
181 - + head * stride_cache_head
182 - + (safe_slot % block_size) * stride_cache_token
183 - + dim * stride_cache_dim
184 - )
185 - if USE_PDL:
186 - tl.extra.cuda.gdc_wait()
187 - packed = _wait_pairs(input_u32, offsets, mask)
188 -
189 - lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
190 - hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
191 - reduced = tl.interleave(
192 - tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16),
193 - tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16),
194 - )
195 - tl.store(
196 - input_u32 + offsets,
197 - tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32),
198 - mask=mask,
199 - )
200 -
201 - tl.store(dst, reduced, mask=valid & channel_mask)
202 - if USE_PDL:
203 - tl.extra.cuda.gdc_launch_dependents()
204 -
205 -
206 -@triton.jit
207 -def _sconv_publish_kernel(
208 - peer_ptrs,
209 - peer_offset_u32,
210 - residual_ptr,
211 - weight_ptr,
212 - cache_ptr,
213 - position_ptr,
214 - sequence_ptr,
215 - slot_ptr,
216 - block_table_ptr,
217 - stride_residual_t,
218 - stride_cache_block,
219 - stride_cache_head,
220 - stride_cache_token,
221 - stride_cache_dim,
222 - stride_block_table_r,
223 - max_blocks,
224 - block_size,
225 - C: tl.constexpr,
226 - CS: tl.constexpr,
227 - CS_P2: tl.constexpr,
228 - SPLITS: tl.constexpr,
229 - HEAD_SIZE: tl.constexpr,
230 - CACHE_OFFSET: tl.constexpr,
231 - RANK: tl.constexpr,
232 - WORLD: tl.constexpr,
233 - WINDOW: tl.constexpr,
234 - USE_PDL: tl.constexpr,
235 - launch_pdl: tl.constexpr,
236 -):
237 - """Compute this rank's shard, then publish it to every rank."""
238 - tl.static_assert(WINDOW == 4)
239 - token = tl.program_id(0).to(tl.int64)
240 - split = tl.program_id(1).to(tl.int64)
241 - CSS: tl.constexpr = CS // SPLITS
242 - channel = tl.arange(0, CS_P2)
243 - channel_mask = channel < CSS
244 - global_channel = split * CSS + channel
245 - head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1)
246 - dim = CACHE_OFFSET + global_channel % HEAD_SIZE
247 - slot = tl.load(slot_ptr + token)
248 - valid = slot >= 0
249 - position = tl.load(position_ptr + token)
250 - sequence = tl.load(sequence_ptr + token)
251 - safe_slot = tl.maximum(slot, 0).to(tl.int64)
252 - own_ptr = (
253 - cache_ptr
254 - + (safe_slot // block_size) * stride_cache_block
255 - + head * stride_cache_head
256 - + (safe_slot % block_size) * stride_cache_token
257 - + dim * stride_cache_dim
258 - )
259 - # These loads were made visible before reduce/insert could signal us and
260 - # are independent of its cache store. Hoisting them gives PDL useful work
261 - # to overlap while retaining the acquire before every cache read.
262 - residual = tl.load(
263 - residual_ptr + token * stride_residual_t + RANK * CS + global_channel,
264 - mask=channel_mask,
265 - other=0.0,
266 - )
267 - weight0 = tl.load(
268 - weight_ptr + global_channel * WINDOW,
269 - mask=channel_mask,
270 - other=0.0,
271 - ).to(tl.float32)
272 - weight1 = tl.load(
273 - weight_ptr + global_channel * WINDOW + 1,
274 - mask=channel_mask,
275 - other=0.0,
276 - ).to(tl.float32)
277 - weight2 = tl.load(
278 - weight_ptr + global_channel * WINDOW + 2,
279 - mask=channel_mask,
280 - other=0.0,
281 - ).to(tl.float32)
282 - weight3 = tl.load(
283 - weight_ptr + global_channel * WINDOW + 3,
284 - mask=channel_mask,
285 - other=0.0,
286 - ).to(tl.float32)
287 - ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64))
288 - if USE_PDL:
289 - tl.extra.cuda.gdc_wait()
290 - current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0)
291 -
292 - conv = tl.zeros([CS_P2], tl.float32)
293 - for tap_idx in tl.static_range(WINDOW):
294 - source_position = position - (WINDOW - 1) + tap_idx
295 - take = valid & (source_position >= 0)
296 - if tap_idx == WINDOW - 1:
297 - value = tl.where(take, current.to(tl.float32), 0.0)
298 - else:
299 - safe_position = tl.maximum(source_position, 0)
300 - logical_block = tl.minimum(safe_position // block_size, max_blocks - 1)
301 - physical_block = tl.load(
302 - block_table_ptr + sequence * stride_block_table_r + logical_block,
303 - mask=take,
304 - other=0,
305 - ).to(tl.int64)
306 - source_ptr = (
307 - cache_ptr
308 - + physical_block * stride_cache_block
309 - + head * stride_cache_head
310 - + (safe_position % block_size) * stride_cache_token
311 - + dim * stride_cache_dim
312 - )
313 - cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0)
314 - value = tl.where(take, cached.to(tl.float32), 0.0)
315 - if tap_idx == 0:
316 - weight = weight0
317 - elif tap_idx == 1:
318 - weight = weight1
319 - elif tap_idx == 2:
320 - weight = weight2
321 - else:
322 - weight = weight3
323 - conv += value * weight
324 -
325 - # Preserve both bf16 rounding points of the original sublayer.
326 - short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16)
327 - output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to(
328 - tl.bfloat16
329 - )
330 -
331 - pair = tl.arange(0, CS_P2 // 2)
332 - pair_mask = pair < CSS // 2
333 - packed = _pack_bf16_pairs(output)
334 - row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair
335 -
336 - for destination in tl.static_range(WORLD):
337 - base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32))
338 - tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask)
339 - if USE_PDL:
340 - tl.extra.cuda.gdc_launch_dependents()
341 -
342 -
343 -@triton.jit
344 -def _gather_norm_kernel(
345 - output_peer_ptrs,
346 - output_peer_offset_u32,
347 - norm_weight_ptr,
348 - normed_ptr,
349 - residual_out_ptr,
350 - eps,
351 - stride_output_t,
352 - C: tl.constexpr,
353 - C_P2: tl.constexpr,
354 - RANK: tl.constexpr,
355 - HAS_NORM: tl.constexpr,
356 - USE_PDL: tl.constexpr,
357 - launch_pdl: tl.constexpr,
358 -):
359 - """Consume a complete row; one CTA owns all outputs for one token."""
360 - token = tl.program_id(0).to(tl.int64)
361 - pair = tl.arange(0, C_P2 // 2)
362 - pair_mask = pair < C // 2
363 - output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64))
364 - output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32))
365 - output_u32 += output_peer_offset_u32
366 - offsets = token * (C // 2) + pair
367 - channel = tl.arange(0, C_P2)
368 - channel_mask = channel < C
369 - # Gamma is independent of the preceding sconv publication. Keep the
370 - # acquire immediately before polling the Lamport output slots.
371 - if HAS_NORM:
372 - weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0)
373 - if USE_PDL:
374 - tl.extra.cuda.gdc_wait()
375 - packed = _wait_pairs(output_u32, offsets, pair_mask)
376 - row = _unpack_bf16_pairs(packed)
377 - tl.store(
378 - residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask
379 - )
380 - if HAS_NORM:
381 - row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0)
382 - inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps)
383 - tl.store(
384 - normed_ptr + token * stride_output_t + channel,
385 - (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16),
386 - mask=channel_mask,
387 - )
388 - tl.store(
389 - output_u32 + offsets,
390 - tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32),
391 - mask=pair_mask,
392 - )
393 - if USE_PDL:
394 - tl.extra.cuda.gdc_launch_dependents()
395 -
396 -
397 -@triton.jit
398 -def _validate_lamport_init_kernel(
399 - input_peer_ptrs,
400 - output_peer_ptrs,
401 - bad_ptr,
402 - num_pairs,
403 - RANK: tl.constexpr,
404 - BLOCK: tl.constexpr,
405 -):
406 - """Validate both complete local allocations through their fabric pointers."""
407 - offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
408 - mask = offsets < num_pairs
409 - ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64))
410 - ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64))
411 - local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32))
412 - local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32))
413 - value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR)
414 - value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR)
415 - bad = tl.max(
416 - tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0)
417 - )
418 - if bad != 0:
419 - tl.atomic_max(bad_ptr, 1)
420 -
421 -
422 -class LamportRSConv:
423 - """Persistent symmetric buffers for one TP group."""
424 -
425 - def _initialize_lamport_buffers(self) -> None:
426 - """Arm every slot and collectively verify the exact sentinel bits.
427 -
428 - Do not replace this with a zero-fill: Lamport readiness distinguishes
429 - the bf16 bit pattern for -0.0 (0x8000) from every published payload.
430 - The validation is deliberately collective so one rank cannot enter the
431 - first polling kernel while another rank still has an unarmed buffer.
432 - """
433 - if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous():
434 - raise RuntimeError("Lamport symmetric buffers must be contiguous")
435 - if self.buf_in.numel() % 2 or self.buf_out.numel() % 2:
436 - raise RuntimeError("Lamport buffers must contain whole uint32 pairs")
437 -
438 - # Fill through int16 so each bf16 lane receives the exact -0.0 bits.
439 - # This covers all three generations and all max-token slots, including
440 - # slots that a smaller first invocation does not touch.
441 - self.buf_in.view(torch.int16).fill_(-0x8000)
442 - self.buf_out.view(torch.int16).fill_(-0x8000)
443 - torch.accelerator.synchronize(self.device)
444 - self.tp.barrier()
445 -
446 - # Scan the complete local allocations. Since every rank performs the
447 - # scan and participates in MAX, success proves every symmetric backing
448 - # allocation was armed before any rank is allowed to use generation 0.
449 - bad = torch.logical_or(
450 - self.buf_in.view(torch.int16).ne(-0x8000).any(),
451 - self.buf_out.view(torch.int16).ne(-0x8000).any(),
452 - ).to(dtype=torch.float32)
453 - bad = self.tp.all_reduce(bad)
454 - if int(bad.item()) != 0:
455 - raise RuntimeError("Lamport sentinel initialization failed on a TP rank")
456 - self.tp.barrier()
457 -
458 - def _initialize_mnnvl_buffers(self) -> None:
459 - """Initialize and validate FlashInfer fabric-mapped Lamport storage."""
460 - self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16)
461 - self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16)
462 - torch.accelerator.synchronize(self.device)
463 - self.tp.barrier()
464 -
465 - bad = torch.zeros((), dtype=torch.int32, device=self.device)
466 - num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2
467 - _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)](
468 - self.input_peer_ptrs,
469 - self.output_peer_ptrs,
470 - bad,
471 - num_pairs,
472 - RANK=self.rank,
473 - BLOCK=256,
474 - num_warps=4,
475 - )
476 - bad = self.tp.all_reduce(bad.to(torch.float32))
477 - if int(bad.item()) != 0:
478 - raise RuntimeError("MNNVL Lamport sentinel initialization failed")
479 - self.tp.barrier()
480 -
481 - def __init__(
482 - self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS
483 - ) -> None:
484 - import torch.distributed._symmetric_memory as symm_mem
485 -
486 - tp = get_tp_group()
487 - self.tp = tp
488 - self.group = tp.device_group
489 - self.world_size = tp.world_size
490 - self.rank = tp.rank_in_group
491 - self.device = torch.device(tp.device)
492 - if self.world_size not in (2, 4, 8):
493 - raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}")
494 - if hidden_size % (2 * self.world_size) != 0:
495 - raise ValueError("hidden size must produce an even shard on every rank")
496 - if window_size != 4:
497 - raise ValueError(f"short-conv window size must be 4, got {window_size}")
498 - if max_tokens < 1 or max_tokens > _MAX_TOKENS:
499 - raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]")
500 -
501 - is_cross_node = not all(in_the_same_node_as(tp.cpu_group))
502 - self.hidden_size = hidden_size
503 - self.window_size = window_size
504 - self.max_tokens = max_tokens
505 - self.shard_size = hidden_size // self.world_size
506 - # Three generations follow FlashInfer's Lamport layout. A generation
507 - # is reused only after two intervening collective calls.
508 - self.num_buffers = 3
509 - self.input_generation_bytes = max_tokens * hidden_size * 2
510 - self.output_generation_bytes = max_tokens * hidden_size * 2
511 - self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9
512 - if is_cross_node:
513 - try:
514 - from flashinfer.comm.mnnvl import (
515 - McastGPUBuffer,
516 - TorchDistBackend,
517 - is_mnnvl_fabric_supported,
518 - )
519 - except ImportError as error:
520 - raise RuntimeError(
521 - "cross-node TP requires FlashInfer MNNVL support"
522 - ) from error
523 -
524 - local_supported = int(
525 - is_mnnvl_fabric_supported(torch.accelerator.current_device_index())
526 - )
527 - unsupported = torch.tensor(
528 - 1 - local_supported, dtype=torch.float32, device=self.device
529 - )
530 - unsupported = tp.all_reduce(unsupported)
531 - if int(unsupported.item()) != 0:
532 - raise RuntimeError("cross-node TP is supported only on MNNVL fabric")
533 -
534 - comm_backend = TorchDistBackend(self.group)
535 - allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2
536 - self._mnnvl_input_handle = McastGPUBuffer(
537 - allocation_bytes,
538 - self.world_size,
539 - self.rank,
540 - self.device,
541 - comm_backend,
542 - )
543 - self._mnnvl_output_handle = McastGPUBuffer(
544 - allocation_bytes,
545 - self.world_size,
546 - self.rank,
547 - self.device,
548 - comm_backend,
549 - )
550 - self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev()
551 - self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev()
552 - self._initialize_mnnvl_buffers()
553 - logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers")
554 - else:
555 - self.buf_in = symm_mem.empty(
556 - self.num_buffers,
557 - max_tokens,
558 - self.world_size,
559 - self.shard_size,
560 - dtype=torch.bfloat16,
561 - device=self.device,
562 - )
563 - self.buf_out = symm_mem.empty(
564 - self.num_buffers,
565 - max_tokens,
566 - hidden_size,
567 - dtype=torch.bfloat16,
568 - device=self.device,
569 - )
570 - group_name = self.group.group_name
571 - input_handle = symm_mem.rendezvous(self.buf_in, group_name)
572 - output_handle = symm_mem.rendezvous(self.buf_out, group_name)
573 - self.input_peer_ptrs = input_handle.buffer_ptrs_dev
574 - self.output_peer_ptrs = output_handle.buffer_ptrs_dev
575 - self._input_handle = input_handle
576 - self._output_handle = output_handle
577 - self._initialize_lamport_buffers()
578 - self.generation = 0
579 -
580 - def usable(self, num_tokens: int) -> bool:
581 - return 0 < num_tokens <= self.max_tokens
582 -
583 - def rs_sconv_ag_add_norm(
584 - self,
585 - input_tensor: torch.Tensor,
586 - residual: torch.Tensor,
587 - conv_weight: torch.Tensor,
588 - norm_weight: torch.Tensor | None,
589 - eps: float,
590 - cache: torch.Tensor,
591 - positions: torch.Tensor,
592 - block_table: torch.Tensor,
593 - seq_idx: torch.Tensor,
594 - slot_mapping: torch.Tensor,
595 - off_s: int,
596 - ws: int,
597 - block_size: int,
598 - ) -> tuple[torch.Tensor | None, torch.Tensor]:
599 - """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``."""
600 - tokens, hidden_size = residual.shape
601 - if not self.usable(tokens):
602 - raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]")
603 - if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16:
604 - raise ValueError("residual must be bf16 [T, 6144]")
605 - if (
606 - input_tensor.shape != residual.shape
607 - or input_tensor.dtype != torch.bfloat16
608 - or input_tensor.stride(1) != 1
609 - ):
610 - raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]")
611 - shard_size = hidden_size // self.world_size
612 - if conv_weight.shape != (shard_size, self.window_size):
613 - raise ValueError(
614 - f"conv_weight must have shape [{shard_size}, {self.window_size}]"
615 - )
616 - if (
617 - conv_weight.dtype != torch.bfloat16
618 - or conv_weight.stride(0) != self.window_size
619 - ):
620 - raise ValueError("conv_weight must be contiguous bf16")
621 - if norm_weight is not None and (
622 - norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16
623 - ):
624 - raise ValueError("norm_weight must be bf16 [6144] or None")
625 - if cache.dtype != torch.bfloat16 or cache.ndim != 4:
626 - raise ValueError("cache must be a 4-D bf16 tensor")
627 - if shard_size % ws != 0 or cache.shape[1] != shard_size // ws:
628 - raise ValueError("cache head layout is inconsistent with ws")
629 - if off_s < 0 or off_s + ws > cache.shape[3]:
630 - raise ValueError("cache channel offset is out of bounds")
631 -
632 - index = self.generation
633 - input_offset = index * self.input_generation_bytes // 4
634 - output_offset = index * self.output_generation_bytes // 4
635 - normed = torch.empty_like(residual) if norm_weight is not None else None
636 - residual_out = torch.empty_like(residual)
637 - phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1
638 - phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits)
639 - phase_grid = (tokens, phase_splits)
640 - # Wide CTAs win before the grid saturates; smaller CTAs reduce register
641 - # pressure once high-throughput batches provide enough parallelism.
642 - if 128 <= tokens <= 2048:
643 - phase_warps = 16
644 - elif tokens > 2048:
645 - phase_warps = 8
646 - else:
647 - phase_warps = 4
648 - gather_warps = 4 if tokens >= 256 else 8
649 - _publish_input_kernel[phase_grid](
650 - input_tensor,
651 - self.input_peer_ptrs,
652 - input_offset,
653 - input_tensor.stride(0),
654 - C=hidden_size,
655 - CS=shard_size,
656 - CS_P2=phase_tile_p2,
657 - SPLITS=phase_splits,
658 - RANK=self.rank,
659 - WORLD=self.world_size,
660 - USE_PDL=self.use_pdl,
661 - launch_pdl=self.use_pdl,
662 - num_warps=phase_warps,
663 - )
664 - _reduce_insert_kernel[phase_grid](
665 - self.input_peer_ptrs,
666 - input_offset,
667 - cache,
668 - slot_mapping,
669 - cache.stride(0),
670 - cache.stride(1),
671 - cache.stride(2),
672 - cache.stride(3),
673 - block_size,
674 - C=hidden_size,
675 - CS=shard_size,
676 - CS_P2=phase_tile_p2,
677 - SPLITS=phase_splits,
678 - HEAD_SIZE=ws,
679 - CACHE_OFFSET=off_s,
680 - RANK=self.rank,
681 - WORLD=self.world_size,
682 - USE_PDL=self.use_pdl,
683 - launch_pdl=self.use_pdl,
684 - num_warps=phase_warps,
685 - )
686 - _sconv_publish_kernel[phase_grid](
687 - self.output_peer_ptrs,
688 - output_offset,
689 - residual,
690 - conv_weight,
691 - cache,
692 - positions,
693 - seq_idx,
694 - slot_mapping,
695 - block_table,
696 - residual.stride(0),
697 - cache.stride(0),
698 - cache.stride(1),
699 - cache.stride(2),
700 - cache.stride(3),
701 - block_table.stride(0),
702 - block_table.shape[1],
703 - block_size,
704 - C=hidden_size,
705 - CS=shard_size,
706 - CS_P2=phase_tile_p2,
707 - SPLITS=phase_splits,
708 - HEAD_SIZE=ws,
709 - CACHE_OFFSET=off_s,
710 - RANK=self.rank,
711 - WORLD=self.world_size,
712 - WINDOW=self.window_size,
713 - USE_PDL=self.use_pdl,
714 - launch_pdl=self.use_pdl,
715 - num_warps=phase_warps,
716 - )
717 - _gather_norm_kernel[(tokens,)](
718 - self.output_peer_ptrs,
719 - output_offset,
720 - norm_weight if norm_weight is not None else residual,
721 - normed if normed is not None else residual_out,
722 - residual_out,
723 - eps,
724 - residual.stride(0),
725 - C=hidden_size,
726 - C_P2=triton.next_power_of_2(hidden_size),
727 - RANK=self.rank,
728 - HAS_NORM=norm_weight is not None,
729 - USE_PDL=self.use_pdl,
730 - launch_pdl=self.use_pdl,
731 - num_warps=gather_warps,
732 - )
733 - self.generation = (index + 1) % self.num_buffers
734 - return normed, residual_out
735 -
736 -
737 -_STATE: LamportRSConv | None = None
738 -_STATE_FAILED = False
739 -
740 -
741 -def initialize_lamport_rs_conv(
742 - hidden_size: int, window_size: int, max_num_batched_tokens: int
743 -) -> None:
744 - """Collectively initialize the TP-group state during model construction."""
745 - global _STATE, _STATE_FAILED
746 - if _STATE is not None:
747 - if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size:
748 - raise RuntimeError("all Lamport users must share hidden and window sizes")
749 - return
750 - if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0":
751 - return
752 - try:
753 - max_tokens = min(_MAX_TOKENS, max_num_batched_tokens)
754 - _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens)
755 - except Exception:
756 - _STATE_FAILED = True
757 - logger.exception("fused collective unavailable; use the NCCL fallback")
758 -
759 -
760 -def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None:
761 - """Return the state initialized with the model, or ``None`` for fallback."""
762 - if _STATE is not None and (
763 - _STATE.hidden_size != hidden_size or _STATE.window_size != window_size
764 - ):
765 - raise RuntimeError("all Lamport users must share hidden and window sizes")
766 - return _STATE
vllm/models/inkling/amd/ops/mm_towers.py deleted
-190
@@ -1,190 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Fused CUDA kernels for the Inkling vision/audio towers.
4 -
5 -Both kernels keep the reference paths' fp32 accumulation and per-op bf16
6 -rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently
7 -bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order
8 -(real-checkpoint-weight cosine vs reference > 0.9999998).
9 -"""
10 -
11 -from __future__ import annotations
12 -
13 -import torch
14 -
15 -from vllm.triton_utils import tl, tldevice, triton
16 -
17 -from .norm import _get_num_warps_from_block_size
18 -
19 -
20 -@triton.jit
21 -def _dmel_embed_sum_norm_kernel(
22 - idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB))
23 - w_ptr, # [NB * VOCAB, D] bf16 embedding table
24 - norm_w_ptr, # [D] (unused if HAS_NORM=False)
25 - out_ptr, # [T, D] bf16
26 - eps,
27 - D,
28 - stride_idx_t,
29 - NB: tl.constexpr,
30 - VOCAB: tl.constexpr,
31 - D_P2: tl.constexpr,
32 - HAS_NORM: tl.constexpr,
33 -):
34 - t = tl.program_id(0).to(tl.int64)
35 - offs = tl.arange(0, D_P2)
36 - mask = offs < D
37 - # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)),
38 - # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()).
39 - acc = tl.zeros([D_P2], dtype=tl.float32)
40 - for b in tl.static_range(NB):
41 - v = tl.load(idx_ptr + t * stride_idx_t + b)
42 - row = (b * VOCAB + v).to(tl.int64)
43 - acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32)
44 - h = acc.to(tl.bfloat16)
45 - if HAS_NORM:
46 - # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded
47 - # bf16 multiply with the bf16 weight.
48 - x32 = h.to(tl.float32)
49 - var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D
50 - xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16)
51 - w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0)
52 - h = xn * w
53 - tl.store(out_ptr + t * D + offs, h, mask=mask)
54 -
55 -
56 -def dmel_embed_sum_norm(
57 - dmel_idx: torch.Tensor, # [T, NB] int32
58 - weight: torch.Tensor, # [NB * VOCAB, D] bf16
59 - norm_weight: torch.Tensor | None,
60 - eps: float,
61 -) -> torch.Tensor:
62 - """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no
63 - [T, NB, D] intermediate)."""
64 - T, nb = dmel_idx.shape
65 - D = weight.shape[1]
66 - assert weight.shape[0] % nb == 0
67 - vocab = weight.shape[0] // nb
68 - out = torch.empty((T, D), dtype=weight.dtype, device=weight.device)
69 - if T == 0:
70 - return out
71 - d_p2 = triton.next_power_of_2(D)
72 - _dmel_embed_sum_norm_kernel[(T,)](
73 - dmel_idx,
74 - weight,
75 - norm_weight if norm_weight is not None else weight,
76 - out,
77 - eps,
78 - D,
79 - dmel_idx.stride(0),
80 - NB=nb,
81 - VOCAB=vocab,
82 - D_P2=d_p2,
83 - HAS_NORM=norm_weight is not None,
84 - # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4%
85 - # at large T (the 80-row gather chain is latency- not lane-bound).
86 - num_warps=8,
87 - )
88 - return out
89 -
90 -
91 -@triton.jit
92 -def _rmsnorm_gelu_kernel(
93 - x_ptr, # [R, D] bf16
94 - w_ptr, # [D]
95 - out_ptr, # [R, D] bf16 (or the folded layout when FOLD)
96 - eps,
97 - R,
98 - D,
99 - D_P2: tl.constexpr,
100 - BLOCK_M: tl.constexpr, # rows per block (>1 for small D)
101 - HAS_GELU: tl.constexpr,
102 - FOLD: tl.constexpr,
103 - # fold geometry: input rows index [N, T, H, W]; the store scatters each
104 - # row to (out_row, slot) of fold_timespace_to_depth's output layout.
105 - FT: tl.constexpr,
106 - FH: tl.constexpr,
107 - FW: tl.constexpr,
108 - TF: tl.constexpr,
109 - HF: tl.constexpr,
110 -):
111 - pid = tl.program_id(0).to(tl.int64)
112 - rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)
113 - rmask = rows < R
114 - offs = tl.arange(0, D_P2)
115 - mask = rmask[:, None] & (offs < D)[None, :]
116 - x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to(
117 - tl.float32
118 - )
119 - var = tl.sum(x32 * x32, axis=1) / D
120 - xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16)
121 - w = tl.load(w_ptr + offs, mask=offs < D, other=0.0)
122 - h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm
123 - if HAS_GELU:
124 - # Exact (erf) GELU on the bf16-rounded norm output, fp32 math,
125 - # matching F.gelu's opmath on a bf16 tensor.
126 - g32 = h.to(tl.float32)
127 - h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16)
128 - if FOLD:
129 - # Store directly into the next layer's folded layout (a pure
130 - # permutation — replaces the separate fold copy pass).
131 - t = (rows // (FH * FW)) % FT
132 - hh = (rows // FW) % FH
133 - ww = rows % FW
134 - n = rows // (FT * FH * FW)
135 - slot = ((t % TF) * HF + hh % HF) * HF + ww % HF
136 - out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * (
137 - FW // HF
138 - ) + ww // HF
139 - base = (out_row * (TF * HF * HF) + slot) * D
140 - tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask)
141 - else:
142 - tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask)
143 -
144 -
145 -def rmsnorm_gelu(
146 - x: torch.Tensor, # [..., D] bf16 contiguous
147 - weight: torch.Tensor,
148 - eps: float,
149 - gelu: bool = True,
150 - fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold
151 -) -> torch.Tensor:
152 - """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block
153 - when D is small. With ``fold``, x must be [N, T, H, W, D] and the output
154 - comes back as ``fold_timespace_to_depth(result, *fold)``."""
155 - D = x.shape[-1]
156 - flat = x.reshape(-1, D)
157 - assert flat.stride(1) == 1 and flat.stride(0) == D
158 - R = flat.shape[0]
159 - if fold is None:
160 - out = torch.empty_like(flat)
161 - ft = fh = fw = tf = hf = 1
162 - out_shape = x.shape
163 - else:
164 - tf, hf = fold
165 - N, ft, fh, fw, _ = x.shape
166 - out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D)
167 - out = torch.empty(out_shape, dtype=x.dtype, device=x.device)
168 - if R == 0:
169 - return out.reshape(out_shape)
170 - d_p2 = triton.next_power_of_2(D)
171 - block_m = max(1, 4096 // d_p2)
172 - _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)](
173 - flat,
174 - weight,
175 - out,
176 - eps,
177 - R,
178 - D,
179 - D_P2=d_p2,
180 - BLOCK_M=block_m,
181 - HAS_GELU=gelu,
182 - FOLD=fold is not None,
183 - FT=ft,
184 - FH=fh,
185 - FW=fw,
186 - TF=tf,
187 - HF=hf,
188 - num_warps=_get_num_warps_from_block_size(d_p2 * block_m),
189 - )
190 - return out.reshape(out_shape)
vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py
-31
@@ -587,34 +587,3 @@ def warp_sum_specialized(
587 mask_and_clamp=31,
588 )
589 return value
590 -
591 -
592 -@cute.jit
593 -def block_sum_specialized(
594 - value: Float32,
595 - warp_sums: cute.Tensor,
596 - tidx: Int32,
597 - warps: cutlass.Constexpr[int],
598 - last_warp_lanes: cutlass.Constexpr[int],
599 - last_warp_mask: cutlass.Constexpr[int],
600 -) -> Float32:
601 - """Upstream-equivalent FP32 block reduction."""
602 -
603 - lane = cute.arch.lane_idx()
604 - warp_idx = cute.arch.warp_idx()
605 - value = warp_sum_specialized(
606 - value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask
607 - )
608 - if lane == 0:
609 - warp_sums[warp_idx] = value
610 - cute.arch.barrier()
611 -
612 - block_sum = Float32(0.0)
613 - if warp_idx == 0:
614 - if lane < Int32(warps):
615 - block_sum = warp_sums[lane]
616 - block_sum = cute.arch.warp_reduction_sum(block_sum)
617 - if lane == 0:
618 - warp_sums[0] = block_sum
619 - cute.arch.barrier()
620 - return warp_sums[0]
vllm/v1/attention/ops/triton_fp8_mqa_logits.py deleted
-262
@@ -1,262 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Temporary gfx942 fallback for AITER's fp8_mqa_logits kernel.
4 -
5 -This module vendors AITER's Triton fp8_mqa_logits kernel with the gfx942
6 -tile-size workaround from ROCm/aiter#3257. It is used only while vLLM's
7 -pinned AITER version lacks that fix.
8 -
9 -TODO: Remove this vendored copy once vLLM pins an AITER version that includes
10 -ROCm/aiter#3257 bugfix for gfx942.
11 -"""
12 -
13 -import torch
14 -
15 -from vllm.triton_utils import tl, triton
16 -
17 -# gfx942 (MI300X) has 64 KiB of LDS per CU. We accept the default
18 -# (BLOCK_KV=128, num_stages=2) tile only when *both* of these hold:
19 -#
20 -# 1. Occupancy gate. With waves_per_eu=2 and num_warps=4 we target two
21 -# workgroups co-resident on a CU -> per-WG LDS budget = 32 KiB. Triton
22 -# keeps Q in registers (loop-invariant) and the fp32 scores accumulator
23 -# in VGPRs (heavy VALU), so only the double-buffered KV tile is
24 -# expected to live in LDS. A 0.9 safety factor leaves headroom for any
25 -# LDS overhead the compiler may add.
26 -#
27 -# 2. Hardware ceiling. Defensive upper bound that also counts Q and
28 -# scores against the 64 KiB CU limit, in case a Triton version (older
29 -# or future) decides to spill them to LDS. False positives here only
30 -# shrink the tile; false negatives are JIT-aborts, so we lean
31 -# conservative.
32 -_GFX942_CU_LDS_BYTES = 64 * 1024
33 -_GFX942_PER_WG_LDS_BUDGET_BYTES = _GFX942_CU_LDS_BYTES * 9 // 20 # ~28.8 KiB
34 -
35 -
36 -def _gfx942_default_tile_fits_lds(num_heads: int, head_size: int) -> bool:
37 - """Return True iff (BLOCK_KV=128, num_stages=2) fits in MI300X LDS."""
38 - BLOCK_KV = 128
39 - NUM_STAGES = 2
40 - kv_bytes = head_size * BLOCK_KV * NUM_STAGES
41 - scores_bytes = num_heads * BLOCK_KV * 4
42 - q_bytes = num_heads * head_size
43 - fits_occupancy = kv_bytes < _GFX942_PER_WG_LDS_BUDGET_BYTES
44 - fits_hardware = q_bytes + kv_bytes + scores_bytes <= _GFX942_CU_LDS_BYTES
45 - return fits_occupancy and fits_hardware
46 -
47 -
48 -@triton.jit
49 -def _fp8_mqa_logits_kernel(
50 - Q_ptr, # fp8e4m3 [seq_len, H, D]
51 - KV_ptr, # fp8e4m3 [seq_len_kv, D]
52 - kv_scales_ptr, # fp32 [seq_len_kv]
53 - weights_ptr, # fp32 [seq_len, H]
54 - cu_start_ptr, # int32 [seq_len]
55 - cu_end_ptr, # int32 [seq_len]
56 - logits_ptr, # fp32 [seq_len, seq_len_kv]
57 - seq_len,
58 - seq_len_kv,
59 - NUM_HEADS: tl.constexpr,
60 - HEAD_SIZE: tl.constexpr,
61 - # strides
62 - stride_q_s: tl.int64,
63 - stride_q_h: tl.constexpr,
64 - stride_q_d: tl.constexpr,
65 - stride_kv_s: tl.int64,
66 - stride_kv_d: tl.constexpr,
67 - stride_w_s: tl.int64,
68 - stride_w_h: tl.constexpr,
69 - stride_logits_s: tl.int64,
70 - stride_logits_k: tl.int64,
71 - # block sizes
72 - BLOCK_KV: tl.constexpr,
73 -):
74 - row_id = tl.program_id(0)
75 - # go from larger to smaller in terms of work
76 - # to reduce the tail effect
77 - row_id = tl.num_programs(0) - row_id - 1
78 - tl.assume(row_id >= 0)
79 - tl.assume(stride_q_s > 0)
80 - tl.assume(stride_q_h > 0)
81 - tl.assume(stride_q_d > 0)
82 - tl.assume(stride_kv_s > 0)
83 - tl.assume(stride_kv_d > 0)
84 - tl.assume(stride_w_s > 0)
85 - tl.assume(stride_w_h > 0)
86 -
87 - logits_row_ptrs = logits_ptr + row_id * stride_logits_s
88 -
89 - h_inds = tl.arange(0, NUM_HEADS)[:, None]
90 - d_inds = tl.arange(0, HEAD_SIZE)
91 -
92 - # load Q[BLOCK_Q, NUM_HEADS, HEAD_SIZE]
93 - q_ptrs = (
94 - Q_ptr + row_id * stride_q_s + h_inds * stride_q_h + d_inds[None, :] * stride_q_d
95 - )
96 -
97 - q_block = tl.load(q_ptrs, cache_modifier=".cg")
98 - w_ptrs = weights_ptr + row_id * stride_w_s + h_inds * stride_w_h
99 - w_block = tl.load(w_ptrs, cache_modifier=".cg").to(tl.float32)
100 -
101 - # Load start/end for each row in this block
102 - start_ind = tl.load(cu_start_ptr + row_id)
103 - end_ind = tl.load(cu_end_ptr + row_id)
104 -
105 - start_ind = tl.maximum(start_ind, 0)
106 - end_ind = tl.minimum(end_ind, seq_len_kv)
107 - shifted_end = end_ind - start_ind
108 - shifted_unmasked_end = shifted_end // BLOCK_KV * BLOCK_KV
109 -
110 - kv_col_offsets = tl.arange(0, BLOCK_KV) + start_ind
111 - kv_ptrs = (
112 - KV_ptr + kv_col_offsets[None, :] * stride_kv_s + d_inds[:, None] * stride_kv_d
113 - )
114 -
115 - kv_scales_ptrs = kv_scales_ptr + kv_col_offsets
116 -
117 - logits_ptrs = logits_row_ptrs + kv_col_offsets * stride_logits_k
118 -
119 - # Loop over KV tiles
120 - for _ in tl.range(0, shifted_unmasked_end, BLOCK_KV):
121 - kv_block = tl.load(kv_ptrs)
122 - kv_scales = tl.load(kv_scales_ptrs)
123 -
124 - # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]
125 - scores = tl.dot(q_block, kv_block, input_precision="ieee")
126 - # Multiply by kv_scales (broadcast along rows)
127 - scores = scores * kv_scales[None, :]
128 - # ReLU
129 - scores = tl.maximum(scores, 0.0)
130 - scores = scores * w_block
131 - # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ]
132 - scores = tl.sum(scores, axis=0)
133 - tl.store(logits_ptrs, scores)
134 -
135 - kv_ptrs += BLOCK_KV * stride_kv_s
136 - kv_scales_ptrs += BLOCK_KV
137 - logits_ptrs += BLOCK_KV * stride_logits_k
138 - kv_col_offsets += BLOCK_KV
139 -
140 - # masked load
141 - kv_col_mask = kv_col_offsets < end_ind
142 - kv_block = tl.load(kv_ptrs, mask=kv_col_mask[None, :], other=0.0)
143 - kv_scales = tl.load(kv_scales_ptrs, mask=kv_col_mask, other=0.0)
144 -
145 - # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]
146 - scores = tl.dot(q_block, kv_block, input_precision="ieee")
147 - # Multiply by kv_scales (broadcast along rows)
148 - scores = scores * kv_scales[None, :]
149 - # ReLU
150 - scores = tl.maximum(scores, 0.0)
151 - scores = scores * w_block
152 - # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ]
153 - scores = tl.sum(scores, axis=0)
154 - # masked store
155 - in_window = (kv_col_offsets >= start_ind) & (kv_col_offsets < end_ind)
156 - tl.store(logits_ptrs, scores, mask=in_window)
157 -
158 -
159 -def fp8_mqa_logits_gfx942(
160 - q: torch.Tensor,
161 - k_fp8: torch.Tensor,
162 - kv_scales: torch.Tensor,
163 - weights: torch.Tensor,
164 - cu_starts: torch.Tensor,
165 - cu_ends: torch.Tensor,
166 -) -> torch.Tensor:
167 - """Compute FP8 MQA logits on MI300X (gfx942) using the vendored kernel.
168 -
169 - Drop-in replacement for ``aiter.ops.triton.attention.fp8_mqa_logits.
170 - fp8_mqa_logits`` on MI300X. Selects ``(BLOCK_KV, num_stages)`` based on
171 - whether the default tile fits within the 64 KiB LDS budget of a gfx942
172 - CU (see module docstring).
173 -
174 - Args:
175 - q: Query tensor of shape ``[M, H, D]``, FP8 dtype.
176 - k_fp8: Key tensor of shape ``[N, D]``, FP8 dtype.
177 - kv_scales: K scales of shape ``[N]`` (or ``[N, 1]`` -- viewed as
178 - ``[N]``), float32.
179 - weights: Per-head weights of shape ``[M, H]``, float32.
180 - cu_starts: Start indices (inclusive) of shape ``[M]``, int32.
181 - cu_ends: End indices (exclusive) of shape ``[M]``, int32.
182 -
183 - Returns:
184 - Logits of shape ``[M, N]``, float32 -- positions outside
185 - ``[cu_starts[i], cu_ends[i])`` for row ``i`` are pre-filled with
186 - ``-inf`` so the caller can run a top-k without masking.
187 - """
188 - seq_len, num_heads, head_size = q.shape
189 - seq_len_kv = k_fp8.shape[0]
190 - assert num_heads & (num_heads - 1) == 0, (
191 - f"num_heads must be a power of two (got {num_heads})"
192 - )
193 - assert head_size & (head_size - 1) == 0, (
194 - f"head_size must be a power of two (got {head_size})"
195 - )
196 -
197 - # The kernel walks ``kv_scales`` as a 1-D contiguous array of size N
198 - # (it indexes by ``kv_scales_ptr + kv_col_offsets``). The vLLM caller
199 - # passes a ``[N, 4]`` uint8 view-cast-to-float32 which lands as
200 - # ``[N, 1]`` contiguous -- byte-identical to ``[N]`` -- but flatten
201 - # explicitly to keep the kernel's pointer arithmetic intent clear.
202 - kv_scales_1d = kv_scales.reshape(-1)
203 -
204 - # Initialise with -inf so positions outside [cu_starts, cu_ends) read
205 - # as ``-inf`` after the masked store path -- this matches AITER's
206 - # ``fp8_mqa_logits`` semantics and is what the top-k consumer expects.
207 - logits = torch.full(
208 - (seq_len, seq_len_kv),
209 - fill_value=-float("inf"),
210 - dtype=torch.float32,
211 - device=q.device,
212 - )
213 -
214 - if _gfx942_default_tile_fits_lds(num_heads, head_size):
215 - block_kv = 128
216 - num_stages = 2
217 - else:
218 - # DSv4 sparse indexer (NUM_HEADS=64, HEAD_SIZE=128) lands here:
219 - # default tile spills past gfx942's 64 KiB LDS budget. (64, 1)
220 - # needs ~33 KiB and clears the per-WG budget with margin.
221 - block_kv = 64
222 - num_stages = 1
223 -
224 - # heuristic for MFMA instruction shape, identical to AITER's choice
225 - matrix_instr_nonkdim = 32
226 - if seq_len <= 1024:
227 - matrix_instr_nonkdim = 16
228 -
229 - stride_q_s, stride_q_h, stride_q_d = q.stride()
230 - stride_kv_s, stride_kv_d = k_fp8.stride()
231 - stride_w_s, stride_w_h = weights.stride()
232 - stride_logits_s, stride_logits_k = logits.stride()
233 -
234 - _fp8_mqa_logits_kernel[(seq_len,)](
235 - Q_ptr=q,
236 - KV_ptr=k_fp8,
237 - kv_scales_ptr=kv_scales_1d,
238 - weights_ptr=weights,
239 - cu_start_ptr=cu_starts,
240 - cu_end_ptr=cu_ends,
241 - logits_ptr=logits,
242 - seq_len=seq_len,
243 - seq_len_kv=seq_len_kv,
244 - NUM_HEADS=num_heads,
245 - HEAD_SIZE=head_size,
246 - stride_q_s=stride_q_s,
247 - stride_q_h=stride_q_h,
248 - stride_q_d=stride_q_d,
249 - stride_kv_s=stride_kv_s,
250 - stride_kv_d=stride_kv_d,
251 - stride_w_s=stride_w_s,
252 - stride_w_h=stride_w_h,
253 - stride_logits_s=stride_logits_s,
254 - stride_logits_k=stride_logits_k,
255 - BLOCK_KV=block_kv,
256 - num_warps=4,
257 - num_stages=num_stages,
258 - waves_per_eu=2,
259 - matrix_instr_nonkdim=matrix_instr_nonkdim,
260 - )
261 -
262 - return logits
vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py
-572
@@ -21,15 +21,11 @@ Stage 2 is reused unchanged from triton_decode_attention.py.
21 """
22
23 import math
24 -from typing import Any
24
25 import torch
26
27 from vllm.platforms import current_platform
28 from vllm.triton_utils import tl, triton
30 -from vllm.v1.attention.ops.triton_decode_attention import _fwd_kernel_stage2
31 -
32 -from .triton_turboquant_decode import _use_fp8_e4b15
29
30 # ROCm prefers num_stages=1 in attention-like kernels to reduce shared-memory
31 # pressure (mirrors the pattern used in triton_decode_attention.py and
@@ -66,387 +62,6 @@ def build_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
62 return lut.reshape(N * N, 2).contiguous()
63
64
69 -# ---------------------------------------------------------------------------
70 -# Stage 1 v2: Grouped-Q + pair LUT + tl.dot + exp2
71 -# ---------------------------------------------------------------------------
72 -
73 -
74 -@triton.jit
75 -def _tq_decode_stage1_v2(
76 - Q_rot_ptr,
77 - KV_cache_ptr,
78 - KV_cache_u16_ptr, # uint16 view — Opt#3 SoA loads
79 - Block_table_ptr,
80 - Seq_lens_ptr,
81 - Centroids_ptr,
82 - Pair_lut_ptr,
83 - Mid_o_ptr,
84 - stride_qb,
85 - stride_qh,
86 - stride_cache_block, # bytes per block (bs*H*slot_aligned)
87 - stride_bt_b,
88 - stride_mid_b,
89 - stride_mid_h,
90 - stride_mid_s,
91 - NUM_Q_HEADS: tl.constexpr,
92 - NUM_KV_HEADS: tl.constexpr,
93 - HEAD_DIM: tl.constexpr,
94 - BLOCK_SIZE: tl.constexpr,
95 - PADDED_SLOT: tl.constexpr,
96 - MAX_NUM_BLOCKS: tl.constexpr,
97 - NUM_KV_SPLITS: tl.constexpr,
98 - KV_GROUP_SIZE: tl.constexpr,
99 - MSE_BITS: tl.constexpr,
100 - MSE_BYTES: tl.constexpr,
101 - VQB: tl.constexpr,
102 - VAL_DATA_BYTES: tl.constexpr,
103 - # Opt#3 SoA layout
104 - KEY_DATA_BYTES: tl.constexpr,
105 - META_REGION_OFFSET: tl.constexpr,
106 - NUM_SOA_FIELDS: tl.constexpr,
107 - SOA_K_NORM: tl.constexpr,
108 - SOA_V_SCALE: tl.constexpr,
109 - SOA_V_ZERO: tl.constexpr,
110 - N_CENTROIDS: tl.constexpr,
111 - ATTN_SCALE: tl.constexpr,
112 - BLOCK_D: tl.constexpr,
113 - TILE_SIZE: tl.constexpr,
114 - BLOCK_M: tl.constexpr,
115 - KEY_FP8: tl.constexpr,
116 - NORM_CORRECTION: tl.constexpr = 0,
117 - FP8_E4B15: tl.constexpr = 0,
118 - USE_PAIR_LUT: tl.constexpr = 0,
119 - USE_BF16_DOT: tl.constexpr = 0,
120 -):
121 - # Grid: (B, Hk, NUM_KV_SPLITS) — one program per KV head group
122 - bid = tl.program_id(0)
123 - kv_hid = tl.program_id(1)
124 - sid = tl.program_id(2)
125 -
126 - seq_len = tl.load(Seq_lens_ptr + bid)
127 - split_len = tl.cdiv(seq_len, NUM_KV_SPLITS)
128 - split_start = split_len * sid
129 - split_end = tl.minimum(split_start + split_len, seq_len)
130 -
131 - if split_start >= split_end:
132 - return
133 -
134 - # Dimension and tile offsets
135 - d_offs = tl.arange(0, BLOCK_D)
136 - d_mask = d_offs < HEAD_DIM
137 - m_offs = tl.arange(0, BLOCK_M)
138 - t_offs = tl.arange(0, TILE_SIZE)
139 -
140 - # Q head indices for this KV group: [kv_hid * GS, kv_hid * GS + GS)
141 - q_head_offs = kv_hid * KV_GROUP_SIZE + m_offs
142 - q_mask = m_offs < KV_GROUP_SIZE
143 -
144 - # Load Q tile: [BLOCK_M, BLOCK_D] float32
145 - q_addrs = bid * stride_qb + q_head_offs[:, None] * stride_qh + d_offs[None, :]
146 - Q = tl.load(
147 - Q_rot_ptr + q_addrs,
148 - mask=q_mask[:, None] & d_mask[None, :],
149 - other=0.0,
150 - )
151 -
152 - # Pre-warm centroids in L1
153 - if not KEY_FP8:
154 - _c_warm = tl.load(Centroids_ptr + tl.arange(0, N_CENTROIDS)) # noqa: F841
155 -
156 - # Loop-invariant bit-extraction constants for MSE keys
157 - if not KEY_FP8:
158 - if MSE_BITS == 4:
159 - pass # 4-bit byte extraction is computed inline in the tile loop
160 - else:
161 - # Generic bit extraction
162 - mse_bit_off = d_offs * MSE_BITS
163 - mse_byte_idx = mse_bit_off // 8
164 - mse_bit_shift = mse_bit_off % 8
165 - mse_mask_val = (1 << MSE_BITS) - 1
166 -
167 - # Loop-invariant bit-extraction for 3-bit values
168 - if VQB == 3:
169 - val_bit_off = d_offs * 3
170 - val_byte_idx = val_bit_off // 8
171 - val_bit_shift = val_bit_off % 8
172 -
173 - # exp2 pre-scaling
174 - RCP_LN2: tl.constexpr = 1.4426950408889634
175 - LN2: tl.constexpr = 0.6931471805599453
176 - QK_SCALE = ATTN_SCALE * RCP_LN2
177 -
178 - # Online softmax accumulators: [BLOCK_M] for max, [BLOCK_M] for L
179 - M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
180 - L = tl.zeros([BLOCK_M], dtype=tl.float32)
181 -
182 - acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32)
183 -
184 - bt_base = bid * stride_bt_b
185 -
186 - # ================================================================
187 - # TILED LOOP: TILE_SIZE KV tokens per iteration
188 - # ================================================================
189 - for start_n in range(split_start, split_end, TILE_SIZE):
190 - kv_offs = start_n + t_offs
191 - kv_mask_1d = kv_offs < split_end
192 -
193 - page_idx = kv_offs // BLOCK_SIZE
194 - page_off = kv_offs % BLOCK_SIZE
195 - block_nums = tl.load(
196 - Block_table_ptr + bt_base + page_idx,
197 - mask=kv_mask_1d,
198 - other=0,
199 - ).to(tl.int64)
200 -
201 - # Opt#3 SoA addressing: data region then per-block SoA metadata.
202 - slot_within_block = page_off.to(tl.int64)
203 - block_base = block_nums * stride_cache_block
204 - DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES
205 - data_bases = (
206 - block_base
207 - + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT)
208 - + tl.cast(kv_hid, tl.int64) * DATA_BYTES_PER_SLOT
209 - )
210 - head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast(
211 - kv_hid, tl.int64
212 - ) * (NUM_SOA_FIELDS * BLOCK_SIZE)
213 - knorm_u16_addrs = (
214 - head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block
215 - )
216 - vscale_u16_addrs = (
217 - head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block
218 - )
219 - vzero_u16_addrs = (
220 - head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block
221 - )
222 -
223 - # ============================================================
224 - # KEY DEQUANT → K_T: [BLOCK_D, TILE_SIZE] for tl.dot(Q, K_T)
225 - # ============================================================
226 - if KEY_FP8:
227 - k_addrs = data_bases[:, None] + d_offs[None, :]
228 - k_raw = tl.load(
229 - KV_cache_ptr + k_addrs,
230 - mask=kv_mask_1d[:, None] & d_mask[None, :],
231 - other=0,
232 - )
233 - if FP8_E4B15:
234 - K_f = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)
235 - else:
236 - K_f = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
237 - # K_f: [TILE_SIZE, BLOCK_D] → transpose for tl.dot
238 - K_T = tl.trans(K_f)
239 - # S: [BLOCK_M, TILE_SIZE] (bf16 dot on ROCm for better dynamic range)
240 - if USE_BF16_DOT:
241 - S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16))
242 - else:
243 - S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16))
244 - else:
245 - # --- MSE dequantization ---
246 - if MSE_BITS == 4 and USE_PAIR_LUT:
247 - # OPT#1 (load-halving): FLUTE §3.2 pair LUT, but with each
248 - # packed byte loaded EXACTLY ONCE instead of twice (the
249 - # previous code built byte_addrs from d_offs // 2 which has
250 - # the pattern [0,0,1,1,2,2,...], duplicating every byte
251 - # load). We load the packed bytes as a [TILE_SIZE, HALF_D]
252 - # tile, decode both nibbles, then issue a single 3-D gather
253 - # against pair_lut which returns (T[lo], T[hi]) in slots
254 - # 0/1. Reshape [TILE_SIZE, HALF_D, 2] -> [TILE_SIZE, BLOCK_D]
255 - # interleaves the pairs into the expected dim order
256 - # (d=2k -> T[lo_k], d=2k+1 -> T[hi_k]).
257 - HALF_D: tl.constexpr = BLOCK_D // 2
258 - half_offs = tl.arange(0, HALF_D)
259 - # Each byte k covers dims (2k, 2k+1); it is in-range iff
260 - # 2k < HEAD_DIM. HEAD_DIM is always even in our layouts,
261 - # so this is equivalent to k < HEAD_DIM // 2.
262 - byte_mask = (half_offs * 2) < HEAD_DIM
263 -
264 - byte_addrs = data_bases[:, None] + half_offs[None, :]
265 - byte_raw = tl.load(
266 - KV_cache_ptr + byte_addrs,
267 - mask=kv_mask_1d[:, None] & byte_mask[None, :],
268 - other=0,
269 - ).to(tl.int32)
270 -
271 - lo_idx = byte_raw & 0xF # [TILE_SIZE, HALF_D]
272 - hi_idx = (byte_raw >> 4) & 0xF # [TILE_SIZE, HALF_D]
273 - pair_key = lo_idx * N_CENTROIDS + hi_idx # [TILE_SIZE, HALF_D]
274 -
275 - # Single 3-D gather: pair_lut[pair_key*2 + {0,1}]
276 - # -> [TILE_SIZE, HALF_D, 2] = (T[lo_k], T[hi_k]) per element.
277 - pair_slot = tl.arange(0, 2)
278 - c_pair = tl.load(
279 - Pair_lut_ptr + pair_key[:, :, None] * 2 + pair_slot[None, None, :],
280 - mask=(kv_mask_1d[:, None, None] & byte_mask[None, :, None]),
281 - other=0.0,
282 - )
283 - # Row-major reshape interleaves pairs:
284 - # out[t, 2k] = c_pair[t, k, 0] = T[lo_k]
285 - # out[t, 2k+1] = c_pair[t, k, 1] = T[hi_k]
286 - c_vals = tl.reshape(c_pair, [TILE_SIZE, BLOCK_D])
287 - elif MSE_BITS == 4:
288 - # 4-bit without pair LUT: exploit byte alignment
289 - half_idx = d_offs // 2
290 - nibble_shift = (d_offs % 2) * 4
291 - mse_addrs = data_bases[:, None] + half_idx[None, :]
292 - mse_raw = tl.load(
293 - KV_cache_ptr + mse_addrs,
294 - mask=kv_mask_1d[:, None] & d_mask[None, :],
295 - other=0,
296 - ).to(tl.int32)
297 - mse_idx = (mse_raw >> nibble_shift[None, :]) & 0xF
298 - c_vals = tl.load(
299 - Centroids_ptr + mse_idx,
300 - mask=kv_mask_1d[:, None] & d_mask[None, :],
301 - other=0.0,
302 - )
303 - else:
304 - # Generic 3-bit path: two byte loads → 16-bit → shift/mask
305 - mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :]
306 - mse_raw0 = tl.load(
307 - KV_cache_ptr + mse_addrs0,
308 - mask=kv_mask_1d[:, None] & d_mask[None, :],
309 - other=0,
310 - ).to(tl.int32)
311 - mse_raw1 = tl.load(
312 - KV_cache_ptr + mse_addrs0 + 1,
313 - mask=kv_mask_1d[:, None] & d_mask[None, :],
314 - other=0,
315 - ).to(tl.int32)
316 - raw16 = mse_raw0 | (mse_raw1 << 8)
317 - mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask_val
318 - c_vals = tl.load(
319 - Centroids_ptr + mse_idx,
320 - mask=kv_mask_1d[:, None] & d_mask[None, :],
321 - other=0.0,
322 - )
323 -
324 - # Opt#1: norm-correction is pre-folded into the stored per-token
325 - # scalar at store time, so no per-tile sum+sqrt+divide here. The
326 - # stored vec_norm carries ||k||/||c_vec||; multiply by c_vals
327 - # below gives the original c_vec/||c_vec|| * ||k||.
328 -
329 - # Opt#3: K-norms from per-block SoA region. Single u16 load per
330 - # token; contiguous across tokens within a block.
331 - norm_u16 = tl.load(
332 - KV_cache_u16_ptr + knorm_u16_addrs, mask=kv_mask_1d, other=0
333 - )
334 - vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32)
335 -
336 - # Reconstruct K: K_recon[t, d] = norm[t] * centroid[t, d]
337 - K_recon = c_vals * vec_norms[:, None]
338 - # K_T: [BLOCK_D, TILE_SIZE]
339 - K_T = tl.trans(K_recon)
340 -
341 - # S: [BLOCK_M, TILE_SIZE] via tensor core MMA
342 - if USE_BF16_DOT:
343 - S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16))
344 - else:
345 - S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16))
346 -
347 - # Mask out-of-range positions
348 - S = tl.where(kv_mask_1d[None, :], S, float("-inf"))
349 -
350 - # ============================================================
351 - # ONLINE SOFTMAX with exp2
352 - # ============================================================
353 - m_j = tl.maximum(M, tl.max(S, axis=1))
354 - m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
355 - P = tl.math.exp2(S - m_j[:, None])
356 - l_j = tl.sum(P, axis=1)
357 - alpha = tl.math.exp2(M - m_j)
358 -
359 - acc = acc * alpha[:, None]
360 - L = L * alpha + l_j
361 - M = m_j
362 -
363 - # ============================================================
364 - # VALUE DEQUANT → V: [TILE_SIZE, BLOCK_D] for tl.dot(P, V)
365 - # Opt#3: V data lives at data_base + KEY_DATA_BYTES; scale/zero
366 - # move to the SoA metadata region (one u16 load each per token).
367 - # ============================================================
368 - val_bases = data_bases + KEY_DATA_BYTES
369 -
370 - if VQB == 3:
371 - val_addrs0 = val_bases[:, None] + val_byte_idx[None, :]
372 - val_raw0 = tl.load(
373 - KV_cache_ptr + val_addrs0,
374 - mask=kv_mask_1d[:, None] & d_mask[None, :],
375 - other=0,
376 - ).to(tl.int32)
377 - val_raw1 = tl.load(
378 - KV_cache_ptr + val_addrs0 + 1,
379 - mask=kv_mask_1d[:, None] & d_mask[None, :],
380 - other=0,
381 - ).to(tl.int32)
382 - raw16 = val_raw0 | (val_raw1 << 8)
383 - v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32)
384 - else: # VQB == 4
385 - # OPT#3 (load-halving, value path): same pattern as OPT#1 for
386 - # keys -- load each byte exactly once as [TILE_SIZE, HALF_D],
387 - # decode both nibbles, then tl.interleave to [TILE_SIZE, BLOCK_D]
388 - # with pattern [v_lo0, v_hi0, v_lo1, v_hi1, ...].
389 - V_HALF_D: tl.constexpr = BLOCK_D // 2
390 - v_half_offs = tl.arange(0, V_HALF_D)
391 - v_byte_mask = (v_half_offs * 2) < HEAD_DIM
392 -
393 - val_addrs = val_bases[:, None] + v_half_offs[None, :]
394 - val_byte = tl.load(
395 - KV_cache_ptr + val_addrs,
396 - mask=kv_mask_1d[:, None] & v_byte_mask[None, :],
397 - other=0,
398 - ).to(tl.int32)
399 - v_lo = (val_byte & 0xF).to(tl.float32) # [T, HALF_D]
400 - v_hi = ((val_byte >> 4) & 0xF).to(tl.float32) # [T, HALF_D]
401 - v_idx = tl.interleave(v_lo, v_hi) # [T, BLOCK_D]
402 -
403 - scale_u16 = tl.load(
404 - KV_cache_u16_ptr + vscale_u16_addrs, mask=kv_mask_1d, other=0
405 - )
406 - zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=kv_mask_1d, other=0)
407 - v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32)
408 - v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32)
409 - V = v_idx * v_scales[:, None] + v_zeros[:, None]
410 -
411 - # P·V accumulation via tensor core MMA
412 - # P: [BLOCK_M, TILE_SIZE], V: [TILE_SIZE, BLOCK_D]
413 - if USE_BF16_DOT:
414 - acc += tl.dot(P.to(tl.bfloat16), V.to(tl.bfloat16))
415 - else:
416 - acc += tl.dot(P.to(tl.float16), V.to(tl.float16))
417 -
418 - # ================================================================
419 - # EPILOGUE: Store per-Q-head partial results for stage2
420 - # ================================================================
421 - # acc: [BLOCK_M, BLOCK_D], L: [BLOCK_M], M: [BLOCK_M]
422 - # Stage2 expects per-Q-head: mid_o[B, Hq, split, D+1]
423 - safe_L = tl.where(L > 0.0, L, 1.0)
424 - one_over_L = 1.0 / safe_L
425 - acc = acc * one_over_L[:, None]
426 -
427 - # Convert M from log2 space back to natural log for stage2 compatibility
428 - lse = M * LN2 + tl.log(safe_L)
429 -
430 - # Store each Q head's result using 2D scatter
431 - # q_head_offs: [BLOCK_M] = kv_hid * KV_GROUP_SIZE + m_offs
432 - out_addrs = (
433 - bid * stride_mid_b
434 - + q_head_offs[:, None] * stride_mid_h
435 - + sid * stride_mid_s
436 - + d_offs[None, :]
437 - )
438 - tl.store(
439 - Mid_o_ptr + out_addrs,
440 - acc,
441 - mask=q_mask[:, None] & d_mask[None, :],
442 - )
443 - # Store LSE per Q head
444 - lse_addrs = (
445 - bid * stride_mid_b + q_head_offs * stride_mid_h + sid * stride_mid_s + HEAD_DIM
446 - )
447 - tl.store(Mid_o_ptr + lse_addrs, lse, mask=q_mask)
448 -
449 -
65 # ---------------------------------------------------------------------------
66 # Launcher v2
67 # ---------------------------------------------------------------------------
@@ -480,190 +95,3 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
95 shows up on a profile, cache by a hash-of-values fingerprint instead.
96 """
97 return build_pair_lut(centroids)
483 -
484 -
485 -def triton_turboquant_decode_attention_v2(
486 - query: torch.Tensor,
487 - kv_cache: torch.Tensor,
488 - block_table: torch.Tensor,
489 - seq_lens: torch.Tensor,
490 - Pi: torch.Tensor,
491 - centroids: torch.Tensor,
492 - scale: float,
493 - mse_bits: int,
494 - key_packed_size: int,
495 - value_quant_bits: int,
496 - value_packed_size: int,
497 - max_seq_len: int = 0, # unused; kept for backward compatibility
498 - key_fp8: bool = False,
499 - norm_correction: bool = False,
500 - PiT: torch.Tensor | None = None,
501 - # Pre-allocated buffers (optional, required for CUDA-graph stability).
502 - mid_o_buf: torch.Tensor | None = None,
503 - output_buf: torch.Tensor | None = None,
504 - lse_buf: torch.Tensor | None = None,
505 - buf_holder: Any = None,
506 - # Fixed split count — MUST be a compile-time constant across iterations
507 - # for CUDA-graph capture/replay to work. Mirrors the v1 launcher contract.
508 - max_num_kv_splits: int = 32,
509 -) -> torch.Tensor:
510 - """Launch optimized TQ decode attention (v2 stage1 + shared stage2).
511 -
512 - Follows the same buffer-reuse + fixed-grid contract as the v1 launcher
513 - (triton_turboquant_decode_attention) so the backend can capture a CUDA
514 - graph across both versions.
515 - """
516 - B, Hq, D = query.shape
517 - Hk = kv_cache.shape[2]
518 - block_size = kv_cache.shape[1]
519 - padded_slot = kv_cache.shape[3]
520 - max_num_blocks = block_table.shape[1]
521 - n_centroids = centroids.shape[0]
522 - kv_group_size = Hq // Hk
523 - device = query.device
524 - del max_seq_len # no longer used: splits is fixed via max_num_kv_splits
525 -
526 - cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size)
527 -
528 - # Opt#3 SoA layout constants (match store-side computation).
529 - key_data_bytes = D if key_fp8 else cfg["mse_bytes"]
530 - data_bytes_per_slot = key_data_bytes + cfg["val_data_bytes"]
531 - meta_region_offset = block_size * Hk * data_bytes_per_slot
532 - num_soa_fields = 2 if key_fp8 else 3
533 - soa_k_norm = 0
534 - soa_v_scale = 0 if key_fp8 else 1
535 - soa_v_zero = 1 if key_fp8 else 2
536 - kv_cache_u16 = kv_cache.view(torch.uint16)
537 -
538 - # Compute q_rot = q @ Pi.T
539 - if key_fp8:
540 - q_rot = query.float().contiguous()
541 - else:
542 - q_float = query.float()
543 - if PiT is None:
544 - PiT = Pi.T.contiguous()
545 - q_rot = (q_float @ PiT).contiguous()
546 -
547 - # BLOCK_M: pad KV_GROUP_SIZE to power of 2, minimum 16 for tensor cores
548 - BLOCK_M = max(16, triton.next_power_of_2(kv_group_size))
549 -
550 - # TILE_SIZE (BLOCK_KV): tokens per inner-loop iteration
551 - TILE_SIZE = 16
552 -
553 - # Fixed split count — must be constant across calls for cudagraph replay.
554 - NUM_KV_SPLITS = max_num_kv_splits
555 -
556 - # --- mid_o buffer reuse (same pattern as v1 launcher) ---
557 - if (
558 - mid_o_buf is not None
559 - and mid_o_buf.shape[0] >= B
560 - and mid_o_buf.shape[2] >= NUM_KV_SPLITS
561 - ):
562 - mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :]
563 - else:
564 - mid_o = torch.empty(
565 - B,
566 - Hq,
567 - NUM_KV_SPLITS,
568 - D + 1,
569 - dtype=torch.float32,
570 - device=device,
571 - )
572 - if buf_holder is not None:
573 - buf_holder._tq_mid_o_buf = mid_o
574 -
575 - fp8_e4b15 = _use_fp8_e4b15(device.index or 0)
576 -
577 - # Build pair LUT for 4-bit MSE (FLUTE §3.2)
578 - use_pair_lut = mse_bits == 4 and not key_fp8
579 - pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids
580 -
581 - # Platform-dependent pipelining depth: ROCm prefers num_stages=1.
582 - stage1_num_stages = 1 if _is_hip else 2
583 -
584 - # --- Stage 1: v2 kernel ---
585 - # Grid over KV heads (not Q heads) — each program handles all Q heads
586 - grid = (B, Hk, NUM_KV_SPLITS)
587 - _tq_decode_stage1_v2[grid](
588 - q_rot,
589 - kv_cache,
590 - kv_cache_u16,
591 - block_table,
592 - seq_lens,
593 - centroids,
594 - pair_lut,
595 - mid_o,
596 - q_rot.stride(0),
597 - q_rot.stride(1),
598 - kv_cache.stride(0),
599 - block_table.stride(0),
600 - mid_o.stride(0),
601 - mid_o.stride(1),
602 - mid_o.stride(2),
603 - NUM_Q_HEADS=Hq,
604 - NUM_KV_HEADS=Hk,
605 - HEAD_DIM=D,
606 - BLOCK_SIZE=block_size,
607 - PADDED_SLOT=padded_slot,
608 - MAX_NUM_BLOCKS=max_num_blocks,
609 - NUM_KV_SPLITS=NUM_KV_SPLITS,
610 - KV_GROUP_SIZE=kv_group_size,
611 - MSE_BITS=mse_bits,
612 - MSE_BYTES=cfg["mse_bytes"],
613 - VQB=value_quant_bits,
614 - VAL_DATA_BYTES=cfg["val_data_bytes"],
615 - KEY_DATA_BYTES=key_data_bytes,
616 - META_REGION_OFFSET=meta_region_offset,
617 - NUM_SOA_FIELDS=num_soa_fields,
618 - SOA_K_NORM=soa_k_norm,
619 - SOA_V_SCALE=soa_v_scale,
620 - SOA_V_ZERO=soa_v_zero,
621 - N_CENTROIDS=n_centroids,
622 - ATTN_SCALE=scale,
623 - BLOCK_D=cfg["BLOCK_D"],
624 - TILE_SIZE=TILE_SIZE,
625 - BLOCK_M=BLOCK_M,
626 - KEY_FP8=1 if key_fp8 else 0,
627 - NORM_CORRECTION=1 if norm_correction else 0,
628 - FP8_E4B15=fp8_e4b15,
629 - USE_PAIR_LUT=1 if use_pair_lut else 0,
630 - USE_BF16_DOT=1 if _is_hip else 0,
631 - num_warps=4,
632 - num_stages=stage1_num_stages,
633 - )
634 -
635 - # --- output / lse buffer reuse (same pattern as v1 launcher) ---
636 - if output_buf is not None and output_buf.shape[0] >= B:
637 - output = output_buf[:B, :Hq, :D]
638 - else:
639 - output = torch.empty(B, Hq, D, dtype=torch.float32, device=device)
640 - if buf_holder is not None:
641 - buf_holder._tq_output_buf = output
642 - if lse_buf is not None and lse_buf.shape[0] >= B:
643 - lse = lse_buf[:B, :Hq]
644 - else:
645 - lse = torch.empty(B, Hq, dtype=torch.float32, device=device)
646 - if buf_holder is not None:
647 - buf_holder._tq_lse_buf = lse
648 -
649 - # --- Stage 2: reduce across KV splits (unchanged) ---
650 - grid2 = (B, Hq)
651 - _fwd_kernel_stage2[grid2](
652 - mid_o,
653 - output,
654 - lse,
655 - seq_lens,
656 - mid_o.stride(0),
657 - mid_o.stride(1),
658 - mid_o.stride(2),
659 - output.stride(0),
660 - output.stride(1),
661 - lse.stride(0),
662 - NUM_KV_SPLITS=NUM_KV_SPLITS,
663 - BLOCK_DV=cfg["BLOCK_D"],
664 - Lv=D,
665 - num_warps=4,
666 - num_stages=2,
667 - )
668 -
669 - return output.to(query.dtype)
vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py
-15
@@ -876,21 +876,6 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
876 return build_pair_lut(centroids)
877
878
879 -# PiT in Q's dtype cached per (PiT-id, dtype) to avoid a host cast on every
880 -# launcher call. Keyed by the fp32 PiT's data_ptr so ownership/lifetime is
881 -# tied to the caller's tensor.
882 -_pit_qdtype_cache: dict = {}
883 -
884 -
885 -def _get_pit_in_query_dtype(PiT: torch.Tensor, qdtype: torch.dtype) -> torch.Tensor:
886 - key = (PiT.data_ptr(), qdtype)
887 - cached = _pit_qdtype_cache.get(key)
888 - if cached is None:
889 - cached = PiT.to(qdtype).contiguous()
890 - _pit_qdtype_cache[key] = cached
891 - return cached
892 -
893 -
879 def triton_turboquant_unified_attention(
880 query: torch.Tensor, # [num_tokens, Hq, D] - fp16/bf16
881 kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8