iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Refactor] Remove dead kernel code (#56845)

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Wentao Ye
2026-09-15 13:17:53 -07:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent df42d112ee
commit 18f8aa0465
8 changed files with 0 additions and 1895 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/mamba/ops/replayssm_config.py
-15
View File
@@ -8,7 +8,6 @@ auto-detected (Blackwell vs not) so call sites need not thread it through.
"""
import functools
from contextlib import contextmanager
from vllm.platforms import current_platform
from vllm.triton_utils import triton
@@ -26,20 +25,6 @@ def _is_blackwell() -> bool:
_overrides: dict[str, tuple] = {}
@contextmanager
def override_replayssm_config(kernel: str, config: tuple):
"""Pin ``kernel``'s launch config for the duration of the context."""
prev = _overrides.get(kernel)
_overrides[kernel] = config
try:
yield
finally:
if prev is None:
_overrides.pop(kernel, None)
else:
_overrides[kernel] = prev
def _dstate_tile(dstate: int, tile: int) -> int:
return max(16, min(tile, triton.next_power_of_2(dstate)))
vllm/models/glm5next/amd/ops/kpool_compress.py
-44
View File
@@ -65,50 +65,6 @@ def _hadamard128(x):
return x * 0.08838834764831845 # 1/sqrt(128)
# Map pool ids to physical cache slots.
def compute_pooled_write_locs(
page_table_64: torch.Tensor,
pool_ids: torch.Tensor,
pool_size: int,
) -> torch.Tensor:
"""Map logical pooled-K ids to physical flat cache slots.
``pool_size`` consecutive tokens share one pool slot that lives at the
*first* token page of each page-group. ``page_table_64`` maps token pages
to physical block ids; we gather the block id of each pool's page-group
and add the in-block pool offset.
"""
assert page_table_64.ndim == 1
pool_ids = pool_ids.to(torch.int64)
block_size = 64
pool_page_group = torch.div(pool_ids, block_size, rounding_mode="floor")
token_page_row = pool_page_group * pool_size
packed_page = page_table_64.index_select(0, token_page_row.to(torch.int64))
return packed_page.to(torch.int64) * block_size + torch.remainder(
pool_ids, block_size
)
def build_pooled_page_table(
page_table: torch.Tensor,
pool_size: int,
) -> torch.Tensor:
"""Build a pool-granular page table by taking every ``pool_size``-th
token-page column (one pool maps to ``pool_size`` token pages).
Uses gather (not strided slicing) so the result is always a fresh
row-major tensor — some downstream kernels require stride(-1) == 1.
"""
block_size = page_table.shape[-1]
assert block_size % pool_size == 0, (
f"pool_size ({pool_size}) must divide page columns ({block_size})"
)
idx = torch.arange(0, block_size, pool_size, device=page_table.device)
return page_table[..., idx].contiguous()
# Fused pool compression and cache write.
vllm/models/inkling/amd/ops/lamport.py
-766
View File
@@ -1,766 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm.
The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``.
Liveness
--------
Large grids are deliberately split at the two communication dependencies:
1. ``_publish_input_kernel`` only publishes rank partials.
2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard.
3. ``_sconv_publish_kernel`` only computes and publishes the local result.
4. ``_gather_norm_kernel`` waits, gathers, and normalizes.
Without PDL, CUDA stream order completes a producer before its consumer. With
PDL, every producer CTA posts all peer stores before triggering its dependent,
and the consumer executes ``gdc_wait`` before polling. A consumer therefore
waits only for stores from a producer that is already running (or complete) on
another GPU. It never waits for another block in its own grid. Consequently
spinning consumers cannot occupy resources needed by any producer, and the
wait-for graph has no cycle. The proof is independent of grid size, block
dispatch order, and occupancy.
For one token, the first three phases use eight independent channel slices to
expose enough CTA parallelism for decode latency. Each slice has exclusive
ownership of its cache and Lamport columns; the gather/RMSNorm phase retains
one CTA per token so no cross-CTA reduction or completion counter is needed.
Immediate buffer reuse (including replay of a captured CUDA graph) is also
safe. Rank R cannot republish input for call n+1 until its gather for n has
finished; that gather waited for owner O's output, which O publishes only
after consuming R's input for n. Likewise, R cannot republish output for n+1
until its reduction for n+1 has observed destination D's input; D publishes
that input only after its gather consumed R's output for n. Thus every prior
read happens-before a same-slot rewrite. Three generations reduce incidental
coupling for ordinary launches, but correctness does not depend on rotation.
The payload itself is the Lamport flag. A 32-bit store publishes two bf16s
atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real
negative zeroes are changed to positive zero before publication. Consumers
use volatile 32-bit loads and restore the sentinel after consuming a slot.
"""
from __future__ import annotations
import os
import torch
from vllm.distributed import get_tp_group
from vllm.distributed.parallel_state import in_the_same_node_as
from vllm.logger import init_logger
from vllm.triton_utils import tl, triton
logger = init_logger(__name__)
_MAX_TOKENS = 16384
_EMPTY_PAIR = tl.constexpr(0x80008000)
@triton.jit
def _pack_bf16_pairs(values):
"""Pack bf16 values into atomic u32 pairs and reserve negative zero."""
lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2]))
lo = lo.to(tl.uint16, bitcast=True)
hi = hi.to(tl.uint16, bitcast=True)
lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32)
hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32)
return lo | (hi << 16)
@triton.jit
def _unpack_bf16_pairs(values):
lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
return tl.interleave(lo, hi)
@triton.jit
def _wait_pairs(ptr, offsets, mask):
values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True)
while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0:
values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True)
return values
@triton.jit
def _publish_input_kernel(
stage_ptr,
peer_ptrs,
peer_offset_u32,
stride_stage_t,
C: tl.constexpr,
CS: tl.constexpr,
CS_P2: tl.constexpr,
SPLITS: tl.constexpr,
RANK: tl.constexpr,
WORLD: tl.constexpr,
USE_PDL: tl.constexpr,
launch_pdl: tl.constexpr,
):
"""Publish this rank's full partial row into every shard owner's slots."""
token = tl.program_id(0).to(tl.int64)
split = tl.program_id(1).to(tl.int64)
CSS: tl.constexpr = CS // SPLITS
pair = tl.arange(0, CS_P2 // 2)
pair_mask = pair < CSS // 2
elem = tl.arange(0, CS_P2)
elem_mask = elem < CSS
ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64))
# Address generation is independent of the preceding stream kernel. The
# acquire remains before stage/peer loads, which may consume its writes.
if USE_PDL:
tl.extra.cuda.gdc_wait()
for owner in tl.static_range(WORLD):
values = tl.load(
stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem,
mask=elem_mask,
other=0.0,
)
packed = _pack_bf16_pairs(values)
base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32))
dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair
tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask)
if USE_PDL:
tl.extra.cuda.gdc_launch_dependents()
@triton.jit
def _reduce_insert_kernel(
input_peer_ptrs,
input_peer_offset_u32,
cache_ptr,
slot_ptr,
stride_cache_block,
stride_cache_head,
stride_cache_token,
stride_cache_dim,
block_size,
C: tl.constexpr,
CS: tl.constexpr,
CS_P2: tl.constexpr,
SPLITS: tl.constexpr,
HEAD_SIZE: tl.constexpr,
CACHE_OFFSET: tl.constexpr,
RANK: tl.constexpr,
WORLD: tl.constexpr,
USE_PDL: tl.constexpr,
launch_pdl: tl.constexpr,
):
"""Consume all partials for this rank and insert the reduced cache row."""
token = tl.program_id(0).to(tl.int64)
split = tl.program_id(1).to(tl.int64)
CSS: tl.constexpr = CS // SPLITS
source = tl.arange(0, WORLD)
pair = tl.arange(0, CS_P2 // 2)
pair_mask = pair < CSS // 2
offsets = (
(token * WORLD + source)[:, None] * (CS // 2)
+ split * (CSS // 2)
+ pair[None, :]
)
mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :]
input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64))
input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32))
input_u32 += input_peer_offset_u32
# Slot metadata and the cache destination do not depend on input publish.
slot = tl.load(slot_ptr + token)
valid = slot >= 0
channel = tl.arange(0, CS_P2)
channel_mask = channel < CSS
global_channel = split * CSS + channel
head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1)
dim = CACHE_OFFSET + global_channel % HEAD_SIZE
safe_slot = tl.maximum(slot, 0).to(tl.int64)
dst = (
cache_ptr
+ (safe_slot // block_size) * stride_cache_block
+ head * stride_cache_head
+ (safe_slot % block_size) * stride_cache_token
+ dim * stride_cache_dim
)
if USE_PDL:
tl.extra.cuda.gdc_wait()
packed = _wait_pairs(input_u32, offsets, mask)
lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
reduced = tl.interleave(
tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16),
tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16),
)
tl.store(
input_u32 + offsets,
tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32),
mask=mask,
)
tl.store(dst, reduced, mask=valid & channel_mask)
if USE_PDL:
tl.extra.cuda.gdc_launch_dependents()
@triton.jit
def _sconv_publish_kernel(
peer_ptrs,
peer_offset_u32,
residual_ptr,
weight_ptr,
cache_ptr,
position_ptr,
sequence_ptr,
slot_ptr,
block_table_ptr,
stride_residual_t,
stride_cache_block,
stride_cache_head,
stride_cache_token,
stride_cache_dim,
stride_block_table_r,
max_blocks,
block_size,
C: tl.constexpr,
CS: tl.constexpr,
CS_P2: tl.constexpr,
SPLITS: tl.constexpr,
HEAD_SIZE: tl.constexpr,
CACHE_OFFSET: tl.constexpr,
RANK: tl.constexpr,
WORLD: tl.constexpr,
WINDOW: tl.constexpr,
USE_PDL: tl.constexpr,
launch_pdl: tl.constexpr,
):
"""Compute this rank's shard, then publish it to every rank."""
tl.static_assert(WINDOW == 4)
token = tl.program_id(0).to(tl.int64)
split = tl.program_id(1).to(tl.int64)
CSS: tl.constexpr = CS // SPLITS
channel = tl.arange(0, CS_P2)
channel_mask = channel < CSS
global_channel = split * CSS + channel
head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1)
dim = CACHE_OFFSET + global_channel % HEAD_SIZE
slot = tl.load(slot_ptr + token)
valid = slot >= 0
position = tl.load(position_ptr + token)
sequence = tl.load(sequence_ptr + token)
safe_slot = tl.maximum(slot, 0).to(tl.int64)
own_ptr = (
cache_ptr
+ (safe_slot // block_size) * stride_cache_block
+ head * stride_cache_head
+ (safe_slot % block_size) * stride_cache_token
+ dim * stride_cache_dim
)
# These loads were made visible before reduce/insert could signal us and
# are independent of its cache store. Hoisting them gives PDL useful work
# to overlap while retaining the acquire before every cache read.
residual = tl.load(
residual_ptr + token * stride_residual_t + RANK * CS + global_channel,
mask=channel_mask,
other=0.0,
)
weight0 = tl.load(
weight_ptr + global_channel * WINDOW,
mask=channel_mask,
other=0.0,
).to(tl.float32)
weight1 = tl.load(
weight_ptr + global_channel * WINDOW + 1,
mask=channel_mask,
other=0.0,
).to(tl.float32)
weight2 = tl.load(
weight_ptr + global_channel * WINDOW + 2,
mask=channel_mask,
other=0.0,
).to(tl.float32)
weight3 = tl.load(
weight_ptr + global_channel * WINDOW + 3,
mask=channel_mask,
other=0.0,
).to(tl.float32)
ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64))
if USE_PDL:
tl.extra.cuda.gdc_wait()
current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0)
conv = tl.zeros([CS_P2], tl.float32)
for tap_idx in tl.static_range(WINDOW):
source_position = position - (WINDOW - 1) + tap_idx
take = valid & (source_position >= 0)
if tap_idx == WINDOW - 1:
value = tl.where(take, current.to(tl.float32), 0.0)
else:
safe_position = tl.maximum(source_position, 0)
logical_block = tl.minimum(safe_position // block_size, max_blocks - 1)
physical_block = tl.load(
block_table_ptr + sequence * stride_block_table_r + logical_block,
mask=take,
other=0,
).to(tl.int64)
source_ptr = (
cache_ptr
+ physical_block * stride_cache_block
+ head * stride_cache_head
+ (safe_position % block_size) * stride_cache_token
+ dim * stride_cache_dim
)
cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0)
value = tl.where(take, cached.to(tl.float32), 0.0)
if tap_idx == 0:
weight = weight0
elif tap_idx == 1:
weight = weight1
elif tap_idx == 2:
weight = weight2
else:
weight = weight3
conv += value * weight
# Preserve both bf16 rounding points of the original sublayer.
short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16)
output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to(
tl.bfloat16
)
pair = tl.arange(0, CS_P2 // 2)
pair_mask = pair < CSS // 2
packed = _pack_bf16_pairs(output)
row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair
for destination in tl.static_range(WORLD):
base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32))
tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask)
if USE_PDL:
tl.extra.cuda.gdc_launch_dependents()
@triton.jit
def _gather_norm_kernel(
output_peer_ptrs,
output_peer_offset_u32,
norm_weight_ptr,
normed_ptr,
residual_out_ptr,
eps,
stride_output_t,
C: tl.constexpr,
C_P2: tl.constexpr,
RANK: tl.constexpr,
HAS_NORM: tl.constexpr,
USE_PDL: tl.constexpr,
launch_pdl: tl.constexpr,
):
"""Consume a complete row; one CTA owns all outputs for one token."""
token = tl.program_id(0).to(tl.int64)
pair = tl.arange(0, C_P2 // 2)
pair_mask = pair < C // 2
output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64))
output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32))
output_u32 += output_peer_offset_u32
offsets = token * (C // 2) + pair
channel = tl.arange(0, C_P2)
channel_mask = channel < C
# Gamma is independent of the preceding sconv publication. Keep the
# acquire immediately before polling the Lamport output slots.
if HAS_NORM:
weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0)
if USE_PDL:
tl.extra.cuda.gdc_wait()
packed = _wait_pairs(output_u32, offsets, pair_mask)
row = _unpack_bf16_pairs(packed)
tl.store(
residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask
)
if HAS_NORM:
row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0)
inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps)
tl.store(
normed_ptr + token * stride_output_t + channel,
(row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16),
mask=channel_mask,
)
tl.store(
output_u32 + offsets,
tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32),
mask=pair_mask,
)
if USE_PDL:
tl.extra.cuda.gdc_launch_dependents()
@triton.jit
def _validate_lamport_init_kernel(
input_peer_ptrs,
output_peer_ptrs,
bad_ptr,
num_pairs,
RANK: tl.constexpr,
BLOCK: tl.constexpr,
):
"""Validate both complete local allocations through their fabric pointers."""
offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
mask = offsets < num_pairs
ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64))
ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64))
local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32))
local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32))
value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR)
value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR)
bad = tl.max(
tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0)
)
if bad != 0:
tl.atomic_max(bad_ptr, 1)
class LamportRSConv:
"""Persistent symmetric buffers for one TP group."""
def _initialize_lamport_buffers(self) -> None:
"""Arm every slot and collectively verify the exact sentinel bits.
Do not replace this with a zero-fill: Lamport readiness distinguishes
the bf16 bit pattern for -0.0 (0x8000) from every published payload.
The validation is deliberately collective so one rank cannot enter the
first polling kernel while another rank still has an unarmed buffer.
"""
if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous():
raise RuntimeError("Lamport symmetric buffers must be contiguous")
if self.buf_in.numel() % 2 or self.buf_out.numel() % 2:
raise RuntimeError("Lamport buffers must contain whole uint32 pairs")
# Fill through int16 so each bf16 lane receives the exact -0.0 bits.
# This covers all three generations and all max-token slots, including
# slots that a smaller first invocation does not touch.
self.buf_in.view(torch.int16).fill_(-0x8000)
self.buf_out.view(torch.int16).fill_(-0x8000)
torch.accelerator.synchronize(self.device)
self.tp.barrier()
# Scan the complete local allocations. Since every rank performs the
# scan and participates in MAX, success proves every symmetric backing
# allocation was armed before any rank is allowed to use generation 0.
bad = torch.logical_or(
self.buf_in.view(torch.int16).ne(-0x8000).any(),
self.buf_out.view(torch.int16).ne(-0x8000).any(),
).to(dtype=torch.float32)
bad = self.tp.all_reduce(bad)
if int(bad.item()) != 0:
raise RuntimeError("Lamport sentinel initialization failed on a TP rank")
self.tp.barrier()
def _initialize_mnnvl_buffers(self) -> None:
"""Initialize and validate FlashInfer fabric-mapped Lamport storage."""
self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16)
self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16)
torch.accelerator.synchronize(self.device)
self.tp.barrier()
bad = torch.zeros((), dtype=torch.int32, device=self.device)
num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2
_validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)](
self.input_peer_ptrs,
self.output_peer_ptrs,
bad,
num_pairs,
RANK=self.rank,
BLOCK=256,
num_warps=4,
)
bad = self.tp.all_reduce(bad.to(torch.float32))
if int(bad.item()) != 0:
raise RuntimeError("MNNVL Lamport sentinel initialization failed")
self.tp.barrier()
def __init__(
self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS
) -> None:
import torch.distributed._symmetric_memory as symm_mem
tp = get_tp_group()
self.tp = tp
self.group = tp.device_group
self.world_size = tp.world_size
self.rank = tp.rank_in_group
self.device = torch.device(tp.device)
if self.world_size not in (2, 4, 8):
raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}")
if hidden_size % (2 * self.world_size) != 0:
raise ValueError("hidden size must produce an even shard on every rank")
if window_size != 4:
raise ValueError(f"short-conv window size must be 4, got {window_size}")
if max_tokens < 1 or max_tokens > _MAX_TOKENS:
raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]")
is_cross_node = not all(in_the_same_node_as(tp.cpu_group))
self.hidden_size = hidden_size
self.window_size = window_size
self.max_tokens = max_tokens
self.shard_size = hidden_size // self.world_size
# Three generations follow FlashInfer's Lamport layout. A generation
# is reused only after two intervening collective calls.
self.num_buffers = 3
self.input_generation_bytes = max_tokens * hidden_size * 2
self.output_generation_bytes = max_tokens * hidden_size * 2
self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9
if is_cross_node:
try:
from flashinfer.comm.mnnvl import (
McastGPUBuffer,
TorchDistBackend,
is_mnnvl_fabric_supported,
)
except ImportError as error:
raise RuntimeError(
"cross-node TP requires FlashInfer MNNVL support"
) from error
local_supported = int(
is_mnnvl_fabric_supported(torch.accelerator.current_device_index())
)
unsupported = torch.tensor(
1 - local_supported, dtype=torch.float32, device=self.device
)
unsupported = tp.all_reduce(unsupported)
if int(unsupported.item()) != 0:
raise RuntimeError("cross-node TP is supported only on MNNVL fabric")
comm_backend = TorchDistBackend(self.group)
allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2
self._mnnvl_input_handle = McastGPUBuffer(
allocation_bytes,
self.world_size,
self.rank,
self.device,
comm_backend,
)
self._mnnvl_output_handle = McastGPUBuffer(
allocation_bytes,
self.world_size,
self.rank,
self.device,
comm_backend,
)
self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev()
self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev()
self._initialize_mnnvl_buffers()
logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers")
else:
self.buf_in = symm_mem.empty(
self.num_buffers,
max_tokens,
self.world_size,
self.shard_size,
dtype=torch.bfloat16,
device=self.device,
)
self.buf_out = symm_mem.empty(
self.num_buffers,
max_tokens,
hidden_size,
dtype=torch.bfloat16,
device=self.device,
)
group_name = self.group.group_name
input_handle = symm_mem.rendezvous(self.buf_in, group_name)
output_handle = symm_mem.rendezvous(self.buf_out, group_name)
self.input_peer_ptrs = input_handle.buffer_ptrs_dev
self.output_peer_ptrs = output_handle.buffer_ptrs_dev
self._input_handle = input_handle
self._output_handle = output_handle
self._initialize_lamport_buffers()
self.generation = 0
def usable(self, num_tokens: int) -> bool:
return 0 < num_tokens <= self.max_tokens
def rs_sconv_ag_add_norm(
self,
input_tensor: torch.Tensor,
residual: torch.Tensor,
conv_weight: torch.Tensor,
norm_weight: torch.Tensor | None,
eps: float,
cache: torch.Tensor,
positions: torch.Tensor,
block_table: torch.Tensor,
seq_idx: torch.Tensor,
slot_mapping: torch.Tensor,
off_s: int,
ws: int,
block_size: int,
) -> tuple[torch.Tensor | None, torch.Tensor]:
"""Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``."""
tokens, hidden_size = residual.shape
if not self.usable(tokens):
raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]")
if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16:
raise ValueError("residual must be bf16 [T, 6144]")
if (
input_tensor.shape != residual.shape
or input_tensor.dtype != torch.bfloat16
or input_tensor.stride(1) != 1
):
raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]")
shard_size = hidden_size // self.world_size
if conv_weight.shape != (shard_size, self.window_size):
raise ValueError(
f"conv_weight must have shape [{shard_size}, {self.window_size}]"
)
if (
conv_weight.dtype != torch.bfloat16
or conv_weight.stride(0) != self.window_size
):
raise ValueError("conv_weight must be contiguous bf16")
if norm_weight is not None and (
norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16
):
raise ValueError("norm_weight must be bf16 [6144] or None")
if cache.dtype != torch.bfloat16 or cache.ndim != 4:
raise ValueError("cache must be a 4-D bf16 tensor")
if shard_size % ws != 0 or cache.shape[1] != shard_size // ws:
raise ValueError("cache head layout is inconsistent with ws")
if off_s < 0 or off_s + ws > cache.shape[3]:
raise ValueError("cache channel offset is out of bounds")
index = self.generation
input_offset = index * self.input_generation_bytes // 4
output_offset = index * self.output_generation_bytes // 4
normed = torch.empty_like(residual) if norm_weight is not None else None
residual_out = torch.empty_like(residual)
phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1
phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits)
phase_grid = (tokens, phase_splits)
# Wide CTAs win before the grid saturates; smaller CTAs reduce register
# pressure once high-throughput batches provide enough parallelism.
if 128 <= tokens <= 2048:
phase_warps = 16
elif tokens > 2048:
phase_warps = 8
else:
phase_warps = 4
gather_warps = 4 if tokens >= 256 else 8
_publish_input_kernel[phase_grid](
input_tensor,
self.input_peer_ptrs,
input_offset,
input_tensor.stride(0),
C=hidden_size,
CS=shard_size,
CS_P2=phase_tile_p2,
SPLITS=phase_splits,
RANK=self.rank,
WORLD=self.world_size,
USE_PDL=self.use_pdl,
launch_pdl=self.use_pdl,
num_warps=phase_warps,
)
_reduce_insert_kernel[phase_grid](
self.input_peer_ptrs,
input_offset,
cache,
slot_mapping,
cache.stride(0),
cache.stride(1),
cache.stride(2),
cache.stride(3),
block_size,
C=hidden_size,
CS=shard_size,
CS_P2=phase_tile_p2,
SPLITS=phase_splits,
HEAD_SIZE=ws,
CACHE_OFFSET=off_s,
RANK=self.rank,
WORLD=self.world_size,
USE_PDL=self.use_pdl,
launch_pdl=self.use_pdl,
num_warps=phase_warps,
)
_sconv_publish_kernel[phase_grid](
self.output_peer_ptrs,
output_offset,
residual,
conv_weight,
cache,
positions,
seq_idx,
slot_mapping,
block_table,
residual.stride(0),
cache.stride(0),
cache.stride(1),
cache.stride(2),
cache.stride(3),
block_table.stride(0),
block_table.shape[1],
block_size,
C=hidden_size,
CS=shard_size,
CS_P2=phase_tile_p2,
SPLITS=phase_splits,
HEAD_SIZE=ws,
CACHE_OFFSET=off_s,
RANK=self.rank,
WORLD=self.world_size,
WINDOW=self.window_size,
USE_PDL=self.use_pdl,
launch_pdl=self.use_pdl,
num_warps=phase_warps,
)
_gather_norm_kernel[(tokens,)](
self.output_peer_ptrs,
output_offset,
norm_weight if norm_weight is not None else residual,
normed if normed is not None else residual_out,
residual_out,
eps,
residual.stride(0),
C=hidden_size,
C_P2=triton.next_power_of_2(hidden_size),
RANK=self.rank,
HAS_NORM=norm_weight is not None,
USE_PDL=self.use_pdl,
launch_pdl=self.use_pdl,
num_warps=gather_warps,
)
self.generation = (index + 1) % self.num_buffers
return normed, residual_out
_STATE: LamportRSConv | None = None
_STATE_FAILED = False
def initialize_lamport_rs_conv(
hidden_size: int, window_size: int, max_num_batched_tokens: int
) -> None:
"""Collectively initialize the TP-group state during model construction."""
global _STATE, _STATE_FAILED
if _STATE is not None:
if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size:
raise RuntimeError("all Lamport users must share hidden and window sizes")
return
if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0":
return
try:
max_tokens = min(_MAX_TOKENS, max_num_batched_tokens)
_STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens)
except Exception:
_STATE_FAILED = True
logger.exception("fused collective unavailable; use the NCCL fallback")
def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None:
"""Return the state initialized with the model, or ``None`` for fallback."""
if _STATE is not None and (
_STATE.hidden_size != hidden_size or _STATE.window_size != window_size
):
raise RuntimeError("all Lamport users must share hidden and window sizes")
return _STATE
vllm/models/inkling/amd/ops/mm_towers.py
-190
View File
@@ -1,190 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Fused CUDA kernels for the Inkling vision/audio towers.
Both kernels keep the reference paths' fp32 accumulation and per-op bf16
rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently
bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order
(real-checkpoint-weight cosine vs reference > 0.9999998).
"""
from __future__ import annotations
import torch
from vllm.triton_utils import tl, tldevice, triton
from .norm import _get_num_warps_from_block_size
@triton.jit
def _dmel_embed_sum_norm_kernel(
idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB))
w_ptr, # [NB * VOCAB, D] bf16 embedding table
norm_w_ptr, # [D] (unused if HAS_NORM=False)
out_ptr, # [T, D] bf16
eps,
D,
stride_idx_t,
NB: tl.constexpr,
VOCAB: tl.constexpr,
D_P2: tl.constexpr,
HAS_NORM: tl.constexpr,
):
t = tl.program_id(0).to(tl.int64)
offs = tl.arange(0, D_P2)
mask = offs < D
# One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)),
# summed in fp32 (matches torch's fp32-accumulated bf16 .sum()).
acc = tl.zeros([D_P2], dtype=tl.float32)
for b in tl.static_range(NB):
v = tl.load(idx_ptr + t * stride_idx_t + b)
row = (b * VOCAB + v).to(tl.int64)
acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32)
h = acc.to(tl.bfloat16)
if HAS_NORM:
# Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded
# bf16 multiply with the bf16 weight.
x32 = h.to(tl.float32)
var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D
xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16)
w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0)
h = xn * w
tl.store(out_ptr + t * D + offs, h, mask=mask)
def dmel_embed_sum_norm(
dmel_idx: torch.Tensor, # [T, NB] int32
weight: torch.Tensor, # [NB * VOCAB, D] bf16
norm_weight: torch.Tensor | None,
eps: float,
) -> torch.Tensor:
"""``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no
[T, NB, D] intermediate)."""
T, nb = dmel_idx.shape
D = weight.shape[1]
assert weight.shape[0] % nb == 0
vocab = weight.shape[0] // nb
out = torch.empty((T, D), dtype=weight.dtype, device=weight.device)
if T == 0:
return out
d_p2 = triton.next_power_of_2(D)
_dmel_embed_sum_norm_kernel[(T,)](
dmel_idx,
weight,
norm_weight if norm_weight is not None else weight,
out,
eps,
D,
dmel_idx.stride(0),
NB=nb,
VOCAB=vocab,
D_P2=d_p2,
HAS_NORM=norm_weight is not None,
# Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4%
# at large T (the 80-row gather chain is latency- not lane-bound).
num_warps=8,
)
return out
@triton.jit
def _rmsnorm_gelu_kernel(
x_ptr, # [R, D] bf16
w_ptr, # [D]
out_ptr, # [R, D] bf16 (or the folded layout when FOLD)
eps,
R,
D,
D_P2: tl.constexpr,
BLOCK_M: tl.constexpr, # rows per block (>1 for small D)
HAS_GELU: tl.constexpr,
FOLD: tl.constexpr,
# fold geometry: input rows index [N, T, H, W]; the store scatters each
# row to (out_row, slot) of fold_timespace_to_depth's output layout.
FT: tl.constexpr,
FH: tl.constexpr,
FW: tl.constexpr,
TF: tl.constexpr,
HF: tl.constexpr,
):
pid = tl.program_id(0).to(tl.int64)
rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)
rmask = rows < R
offs = tl.arange(0, D_P2)
mask = rmask[:, None] & (offs < D)[None, :]
x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to(
tl.float32
)
var = tl.sum(x32 * x32, axis=1) / D
xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16)
w = tl.load(w_ptr + offs, mask=offs < D, other=0.0)
h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm
if HAS_GELU:
# Exact (erf) GELU on the bf16-rounded norm output, fp32 math,
# matching F.gelu's opmath on a bf16 tensor.
g32 = h.to(tl.float32)
h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16)
if FOLD:
# Store directly into the next layer's folded layout (a pure
# permutation — replaces the separate fold copy pass).
t = (rows // (FH * FW)) % FT
hh = (rows // FW) % FH
ww = rows % FW
n = rows // (FT * FH * FW)
slot = ((t % TF) * HF + hh % HF) * HF + ww % HF
out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * (
FW // HF
) + ww // HF
base = (out_row * (TF * HF * HF) + slot) * D
tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask)
else:
tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask)
def rmsnorm_gelu(
x: torch.Tensor, # [..., D] bf16 contiguous
weight: torch.Tensor,
eps: float,
gelu: bool = True,
fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold
) -> torch.Tensor:
"""Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block
when D is small. With ``fold``, x must be [N, T, H, W, D] and the output
comes back as ``fold_timespace_to_depth(result, *fold)``."""
D = x.shape[-1]
flat = x.reshape(-1, D)
assert flat.stride(1) == 1 and flat.stride(0) == D
R = flat.shape[0]
if fold is None:
out = torch.empty_like(flat)
ft = fh = fw = tf = hf = 1
out_shape = x.shape
else:
tf, hf = fold
N, ft, fh, fw, _ = x.shape
out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D)
out = torch.empty(out_shape, dtype=x.dtype, device=x.device)
if R == 0:
return out.reshape(out_shape)
d_p2 = triton.next_power_of_2(D)
block_m = max(1, 4096 // d_p2)
_rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)](
flat,
weight,
out,
eps,
R,
D,
D_P2=d_p2,
BLOCK_M=block_m,
HAS_GELU=gelu,
FOLD=fold is not None,
FT=ft,
FH=fh,
FW=fw,
TF=tf,
HF=hf,
num_warps=_get_num_warps_from_block_size(d_p2 * block_m),
)
return out.reshape(out_shape)
vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py
-31
View File
@@ -587,34 +587,3 @@ def warp_sum_specialized(
mask_and_clamp=31,
)
return value
@cute.jit
def block_sum_specialized(
value: Float32,
warp_sums: cute.Tensor,
tidx: Int32,
warps: cutlass.Constexpr[int],
last_warp_lanes: cutlass.Constexpr[int],
last_warp_mask: cutlass.Constexpr[int],
) -> Float32:
"""Upstream-equivalent FP32 block reduction."""
lane = cute.arch.lane_idx()
warp_idx = cute.arch.warp_idx()
value = warp_sum_specialized(
value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask
)
if lane == 0:
warp_sums[warp_idx] = value
cute.arch.barrier()
block_sum = Float32(0.0)
if warp_idx == 0:
if lane < Int32(warps):
block_sum = warp_sums[lane]
block_sum = cute.arch.warp_reduction_sum(block_sum)
if lane == 0:
warp_sums[0] = block_sum
cute.arch.barrier()
return warp_sums[0]
vllm/v1/attention/ops/triton_fp8_mqa_logits.py
-262
View File
@@ -1,262 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Temporary gfx942 fallback for AITER's fp8_mqa_logits kernel.
This module vendors AITER's Triton fp8_mqa_logits kernel with the gfx942
tile-size workaround from ROCm/aiter#3257. It is used only while vLLM's
pinned AITER version lacks that fix.
TODO: Remove this vendored copy once vLLM pins an AITER version that includes
ROCm/aiter#3257 bugfix for gfx942.
"""
import torch
from vllm.triton_utils import tl, triton
# gfx942 (MI300X) has 64 KiB of LDS per CU. We accept the default
# (BLOCK_KV=128, num_stages=2) tile only when *both* of these hold:
#
# 1. Occupancy gate. With waves_per_eu=2 and num_warps=4 we target two
# workgroups co-resident on a CU -> per-WG LDS budget = 32 KiB. Triton
# keeps Q in registers (loop-invariant) and the fp32 scores accumulator
# in VGPRs (heavy VALU), so only the double-buffered KV tile is
# expected to live in LDS. A 0.9 safety factor leaves headroom for any
# LDS overhead the compiler may add.
#
# 2. Hardware ceiling. Defensive upper bound that also counts Q and
# scores against the 64 KiB CU limit, in case a Triton version (older
# or future) decides to spill them to LDS. False positives here only
# shrink the tile; false negatives are JIT-aborts, so we lean
# conservative.
_GFX942_CU_LDS_BYTES = 64 * 1024
_GFX942_PER_WG_LDS_BUDGET_BYTES = _GFX942_CU_LDS_BYTES * 9 // 20 # ~28.8 KiB
def _gfx942_default_tile_fits_lds(num_heads: int, head_size: int) -> bool:
"""Return True iff (BLOCK_KV=128, num_stages=2) fits in MI300X LDS."""
BLOCK_KV = 128
NUM_STAGES = 2
kv_bytes = head_size * BLOCK_KV * NUM_STAGES
scores_bytes = num_heads * BLOCK_KV * 4
q_bytes = num_heads * head_size
fits_occupancy = kv_bytes < _GFX942_PER_WG_LDS_BUDGET_BYTES
fits_hardware = q_bytes + kv_bytes + scores_bytes <= _GFX942_CU_LDS_BYTES
return fits_occupancy and fits_hardware
@triton.jit
def _fp8_mqa_logits_kernel(
Q_ptr, # fp8e4m3 [seq_len, H, D]
KV_ptr, # fp8e4m3 [seq_len_kv, D]
kv_scales_ptr, # fp32 [seq_len_kv]
weights_ptr, # fp32 [seq_len, H]
cu_start_ptr, # int32 [seq_len]
cu_end_ptr, # int32 [seq_len]
logits_ptr, # fp32 [seq_len, seq_len_kv]
seq_len,
seq_len_kv,
NUM_HEADS: tl.constexpr,
HEAD_SIZE: tl.constexpr,
# strides
stride_q_s: tl.int64,
stride_q_h: tl.constexpr,
stride_q_d: tl.constexpr,
stride_kv_s: tl.int64,
stride_kv_d: tl.constexpr,
stride_w_s: tl.int64,
stride_w_h: tl.constexpr,
stride_logits_s: tl.int64,
stride_logits_k: tl.int64,
# block sizes
BLOCK_KV: tl.constexpr,
):
row_id = tl.program_id(0)
# go from larger to smaller in terms of work
# to reduce the tail effect
row_id = tl.num_programs(0) - row_id - 1
tl.assume(row_id >= 0)
tl.assume(stride_q_s > 0)
tl.assume(stride_q_h > 0)
tl.assume(stride_q_d > 0)
tl.assume(stride_kv_s > 0)
tl.assume(stride_kv_d > 0)
tl.assume(stride_w_s > 0)
tl.assume(stride_w_h > 0)
logits_row_ptrs = logits_ptr + row_id * stride_logits_s
h_inds = tl.arange(0, NUM_HEADS)[:, None]
d_inds = tl.arange(0, HEAD_SIZE)
# load Q[BLOCK_Q, NUM_HEADS, HEAD_SIZE]
q_ptrs = (
Q_ptr + row_id * stride_q_s + h_inds * stride_q_h + d_inds[None, :] * stride_q_d
)
q_block = tl.load(q_ptrs, cache_modifier=".cg")
w_ptrs = weights_ptr + row_id * stride_w_s + h_inds * stride_w_h
w_block = tl.load(w_ptrs, cache_modifier=".cg").to(tl.float32)
# Load start/end for each row in this block
start_ind = tl.load(cu_start_ptr + row_id)
end_ind = tl.load(cu_end_ptr + row_id)
start_ind = tl.maximum(start_ind, 0)
end_ind = tl.minimum(end_ind, seq_len_kv)
shifted_end = end_ind - start_ind
shifted_unmasked_end = shifted_end // BLOCK_KV * BLOCK_KV
kv_col_offsets = tl.arange(0, BLOCK_KV) + start_ind
kv_ptrs = (
KV_ptr + kv_col_offsets[None, :] * stride_kv_s + d_inds[:, None] * stride_kv_d
)
kv_scales_ptrs = kv_scales_ptr + kv_col_offsets
logits_ptrs = logits_row_ptrs + kv_col_offsets * stride_logits_k
# Loop over KV tiles
for _ in tl.range(0, shifted_unmasked_end, BLOCK_KV):
kv_block = tl.load(kv_ptrs)
kv_scales = tl.load(kv_scales_ptrs)
# [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]
scores = tl.dot(q_block, kv_block, input_precision="ieee")
# Multiply by kv_scales (broadcast along rows)
scores = scores * kv_scales[None, :]
# ReLU
scores = tl.maximum(scores, 0.0)
scores = scores * w_block
# [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ]
scores = tl.sum(scores, axis=0)
tl.store(logits_ptrs, scores)
kv_ptrs += BLOCK_KV * stride_kv_s
kv_scales_ptrs += BLOCK_KV
logits_ptrs += BLOCK_KV * stride_logits_k
kv_col_offsets += BLOCK_KV
# masked load
kv_col_mask = kv_col_offsets < end_ind
kv_block = tl.load(kv_ptrs, mask=kv_col_mask[None, :], other=0.0)
kv_scales = tl.load(kv_scales_ptrs, mask=kv_col_mask, other=0.0)
# [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]
scores = tl.dot(q_block, kv_block, input_precision="ieee")
# Multiply by kv_scales (broadcast along rows)
scores = scores * kv_scales[None, :]
# ReLU
scores = tl.maximum(scores, 0.0)
scores = scores * w_block
# [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ]
scores = tl.sum(scores, axis=0)
# masked store
in_window = (kv_col_offsets >= start_ind) & (kv_col_offsets < end_ind)
tl.store(logits_ptrs, scores, mask=in_window)
def fp8_mqa_logits_gfx942(
q: torch.Tensor,
k_fp8: torch.Tensor,
kv_scales: torch.Tensor,
weights: torch.Tensor,
cu_starts: torch.Tensor,
cu_ends: torch.Tensor,
) -> torch.Tensor:
"""Compute FP8 MQA logits on MI300X (gfx942) using the vendored kernel.
Drop-in replacement for ``aiter.ops.triton.attention.fp8_mqa_logits.
fp8_mqa_logits`` on MI300X. Selects ``(BLOCK_KV, num_stages)`` based on
whether the default tile fits within the 64 KiB LDS budget of a gfx942
CU (see module docstring).
Args:
q: Query tensor of shape ``[M, H, D]``, FP8 dtype.
k_fp8: Key tensor of shape ``[N, D]``, FP8 dtype.
kv_scales: K scales of shape ``[N]`` (or ``[N, 1]`` -- viewed as
``[N]``), float32.
weights: Per-head weights of shape ``[M, H]``, float32.
cu_starts: Start indices (inclusive) of shape ``[M]``, int32.
cu_ends: End indices (exclusive) of shape ``[M]``, int32.
Returns:
Logits of shape ``[M, N]``, float32 -- positions outside
``[cu_starts[i], cu_ends[i])`` for row ``i`` are pre-filled with
``-inf`` so the caller can run a top-k without masking.
"""
seq_len, num_heads, head_size = q.shape
seq_len_kv = k_fp8.shape[0]
assert num_heads & (num_heads - 1) == 0, (
f"num_heads must be a power of two (got {num_heads})"
)
assert head_size & (head_size - 1) == 0, (
f"head_size must be a power of two (got {head_size})"
)
# The kernel walks ``kv_scales`` as a 1-D contiguous array of size N
# (it indexes by ``kv_scales_ptr + kv_col_offsets``). The vLLM caller
# passes a ``[N, 4]`` uint8 view-cast-to-float32 which lands as
# ``[N, 1]`` contiguous -- byte-identical to ``[N]`` -- but flatten
# explicitly to keep the kernel's pointer arithmetic intent clear.
kv_scales_1d = kv_scales.reshape(-1)
# Initialise with -inf so positions outside [cu_starts, cu_ends) read
# as ``-inf`` after the masked store path -- this matches AITER's
# ``fp8_mqa_logits`` semantics and is what the top-k consumer expects.
logits = torch.full(
(seq_len, seq_len_kv),
fill_value=-float("inf"),
dtype=torch.float32,
device=q.device,
)
if _gfx942_default_tile_fits_lds(num_heads, head_size):
block_kv = 128
num_stages = 2
else:
# DSv4 sparse indexer (NUM_HEADS=64, HEAD_SIZE=128) lands here:
# default tile spills past gfx942's 64 KiB LDS budget. (64, 1)
# needs ~33 KiB and clears the per-WG budget with margin.
block_kv = 64
num_stages = 1
# heuristic for MFMA instruction shape, identical to AITER's choice
matrix_instr_nonkdim = 32
if seq_len <= 1024:
matrix_instr_nonkdim = 16
stride_q_s, stride_q_h, stride_q_d = q.stride()
stride_kv_s, stride_kv_d = k_fp8.stride()
stride_w_s, stride_w_h = weights.stride()
stride_logits_s, stride_logits_k = logits.stride()
_fp8_mqa_logits_kernel[(seq_len,)](
Q_ptr=q,
KV_ptr=k_fp8,
kv_scales_ptr=kv_scales_1d,
weights_ptr=weights,
cu_start_ptr=cu_starts,
cu_end_ptr=cu_ends,
logits_ptr=logits,
seq_len=seq_len,
seq_len_kv=seq_len_kv,
NUM_HEADS=num_heads,
HEAD_SIZE=head_size,
stride_q_s=stride_q_s,
stride_q_h=stride_q_h,
stride_q_d=stride_q_d,
stride_kv_s=stride_kv_s,
stride_kv_d=stride_kv_d,
stride_w_s=stride_w_s,
stride_w_h=stride_w_h,
stride_logits_s=stride_logits_s,
stride_logits_k=stride_logits_k,
BLOCK_KV=block_kv,
num_warps=4,
num_stages=num_stages,
waves_per_eu=2,
matrix_instr_nonkdim=matrix_instr_nonkdim,
)
return logits
vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py
-572
View File
@@ -21,15 +21,11 @@ Stage 2 is reused unchanged from triton_decode_attention.py.
"""
import math
from typing import Any
import torch
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.v1.attention.ops.triton_decode_attention import _fwd_kernel_stage2
from .triton_turboquant_decode import _use_fp8_e4b15
# ROCm prefers num_stages=1 in attention-like kernels to reduce shared-memory
# pressure (mirrors the pattern used in triton_decode_attention.py and
@@ -66,387 +62,6 @@ def build_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
return lut.reshape(N * N, 2).contiguous()
# ---------------------------------------------------------------------------
# Stage 1 v2: Grouped-Q + pair LUT + tl.dot + exp2
# ---------------------------------------------------------------------------
@triton.jit
def _tq_decode_stage1_v2(
Q_rot_ptr,
KV_cache_ptr,
KV_cache_u16_ptr, # uint16 view — Opt#3 SoA loads
Block_table_ptr,
Seq_lens_ptr,
Centroids_ptr,
Pair_lut_ptr,
Mid_o_ptr,
stride_qb,
stride_qh,
stride_cache_block, # bytes per block (bs*H*slot_aligned)
stride_bt_b,
stride_mid_b,
stride_mid_h,
stride_mid_s,
NUM_Q_HEADS: tl.constexpr,
NUM_KV_HEADS: tl.constexpr,
HEAD_DIM: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
PADDED_SLOT: tl.constexpr,
MAX_NUM_BLOCKS: tl.constexpr,
NUM_KV_SPLITS: tl.constexpr,
KV_GROUP_SIZE: tl.constexpr,
MSE_BITS: tl.constexpr,
MSE_BYTES: tl.constexpr,
VQB: tl.constexpr,
VAL_DATA_BYTES: tl.constexpr,
# Opt#3 SoA layout
KEY_DATA_BYTES: tl.constexpr,
META_REGION_OFFSET: tl.constexpr,
NUM_SOA_FIELDS: tl.constexpr,
SOA_K_NORM: tl.constexpr,
SOA_V_SCALE: tl.constexpr,
SOA_V_ZERO: tl.constexpr,
N_CENTROIDS: tl.constexpr,
ATTN_SCALE: tl.constexpr,
BLOCK_D: tl.constexpr,
TILE_SIZE: tl.constexpr,
BLOCK_M: tl.constexpr,
KEY_FP8: tl.constexpr,
NORM_CORRECTION: tl.constexpr = 0,
FP8_E4B15: tl.constexpr = 0,
USE_PAIR_LUT: tl.constexpr = 0,
USE_BF16_DOT: tl.constexpr = 0,
):
# Grid: (B, Hk, NUM_KV_SPLITS) — one program per KV head group
bid = tl.program_id(0)
kv_hid = tl.program_id(1)
sid = tl.program_id(2)
seq_len = tl.load(Seq_lens_ptr + bid)
split_len = tl.cdiv(seq_len, NUM_KV_SPLITS)
split_start = split_len * sid
split_end = tl.minimum(split_start + split_len, seq_len)
if split_start >= split_end:
return
# Dimension and tile offsets
d_offs = tl.arange(0, BLOCK_D)
d_mask = d_offs < HEAD_DIM
m_offs = tl.arange(0, BLOCK_M)
t_offs = tl.arange(0, TILE_SIZE)
# Q head indices for this KV group: [kv_hid * GS, kv_hid * GS + GS)
q_head_offs = kv_hid * KV_GROUP_SIZE + m_offs
q_mask = m_offs < KV_GROUP_SIZE
# Load Q tile: [BLOCK_M, BLOCK_D] float32
q_addrs = bid * stride_qb + q_head_offs[:, None] * stride_qh + d_offs[None, :]
Q = tl.load(
Q_rot_ptr + q_addrs,
mask=q_mask[:, None] & d_mask[None, :],
other=0.0,
)
# Pre-warm centroids in L1
if not KEY_FP8:
_c_warm = tl.load(Centroids_ptr + tl.arange(0, N_CENTROIDS)) # noqa: F841
# Loop-invariant bit-extraction constants for MSE keys
if not KEY_FP8:
if MSE_BITS == 4:
pass # 4-bit byte extraction is computed inline in the tile loop
else:
# Generic bit extraction
mse_bit_off = d_offs * MSE_BITS
mse_byte_idx = mse_bit_off // 8
mse_bit_shift = mse_bit_off % 8
mse_mask_val = (1 << MSE_BITS) - 1
# Loop-invariant bit-extraction for 3-bit values
if VQB == 3:
val_bit_off = d_offs * 3
val_byte_idx = val_bit_off // 8
val_bit_shift = val_bit_off % 8
# exp2 pre-scaling
RCP_LN2: tl.constexpr = 1.4426950408889634
LN2: tl.constexpr = 0.6931471805599453
QK_SCALE = ATTN_SCALE * RCP_LN2
# Online softmax accumulators: [BLOCK_M] for max, [BLOCK_M] for L
M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
L = tl.zeros([BLOCK_M], dtype=tl.float32)
acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32)
bt_base = bid * stride_bt_b
# ================================================================
# TILED LOOP: TILE_SIZE KV tokens per iteration
# ================================================================
for start_n in range(split_start, split_end, TILE_SIZE):
kv_offs = start_n + t_offs
kv_mask_1d = kv_offs < split_end
page_idx = kv_offs // BLOCK_SIZE
page_off = kv_offs % BLOCK_SIZE
block_nums = tl.load(
Block_table_ptr + bt_base + page_idx,
mask=kv_mask_1d,
other=0,
).to(tl.int64)
# Opt#3 SoA addressing: data region then per-block SoA metadata.
slot_within_block = page_off.to(tl.int64)
block_base = block_nums * stride_cache_block
DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES
data_bases = (
block_base
+ slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT)
+ tl.cast(kv_hid, tl.int64) * DATA_BYTES_PER_SLOT
)
head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast(
kv_hid, tl.int64
) * (NUM_SOA_FIELDS * BLOCK_SIZE)
knorm_u16_addrs = (
head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block
)
vscale_u16_addrs = (
head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block
)
vzero_u16_addrs = (
head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block
)
# ============================================================
# KEY DEQUANT → K_T: [BLOCK_D, TILE_SIZE] for tl.dot(Q, K_T)
# ============================================================
if KEY_FP8:
k_addrs = data_bases[:, None] + d_offs[None, :]
k_raw = tl.load(
KV_cache_ptr + k_addrs,
mask=kv_mask_1d[:, None] & d_mask[None, :],
other=0,
)
if FP8_E4B15:
K_f = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)
else:
K_f = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
# K_f: [TILE_SIZE, BLOCK_D] → transpose for tl.dot
K_T = tl.trans(K_f)
# S: [BLOCK_M, TILE_SIZE] (bf16 dot on ROCm for better dynamic range)
if USE_BF16_DOT:
S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16))
else:
S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16))
else:
# --- MSE dequantization ---
if MSE_BITS == 4 and USE_PAIR_LUT:
# OPT#1 (load-halving): FLUTE §3.2 pair LUT, but with each
# packed byte loaded EXACTLY ONCE instead of twice (the
# previous code built byte_addrs from d_offs // 2 which has
# the pattern [0,0,1,1,2,2,...], duplicating every byte
# load). We load the packed bytes as a [TILE_SIZE, HALF_D]
# tile, decode both nibbles, then issue a single 3-D gather
# against pair_lut which returns (T[lo], T[hi]) in slots
# 0/1. Reshape [TILE_SIZE, HALF_D, 2] -> [TILE_SIZE, BLOCK_D]
# interleaves the pairs into the expected dim order
# (d=2k -> T[lo_k], d=2k+1 -> T[hi_k]).
HALF_D: tl.constexpr = BLOCK_D // 2
half_offs = tl.arange(0, HALF_D)
# Each byte k covers dims (2k, 2k+1); it is in-range iff
# 2k < HEAD_DIM. HEAD_DIM is always even in our layouts,
# so this is equivalent to k < HEAD_DIM // 2.
byte_mask = (half_offs * 2) < HEAD_DIM
byte_addrs = data_bases[:, None] + half_offs[None, :]
byte_raw = tl.load(
KV_cache_ptr + byte_addrs,
mask=kv_mask_1d[:, None] & byte_mask[None, :],
other=0,
).to(tl.int32)
lo_idx = byte_raw & 0xF # [TILE_SIZE, HALF_D]
hi_idx = (byte_raw >> 4) & 0xF # [TILE_SIZE, HALF_D]
pair_key = lo_idx * N_CENTROIDS + hi_idx # [TILE_SIZE, HALF_D]
# Single 3-D gather: pair_lut[pair_key*2 + {0,1}]
# -> [TILE_SIZE, HALF_D, 2] = (T[lo_k], T[hi_k]) per element.
pair_slot = tl.arange(0, 2)
c_pair = tl.load(
Pair_lut_ptr + pair_key[:, :, None] * 2 + pair_slot[None, None, :],
mask=(kv_mask_1d[:, None, None] & byte_mask[None, :, None]),
other=0.0,
)
# Row-major reshape interleaves pairs:
# out[t, 2k] = c_pair[t, k, 0] = T[lo_k]
# out[t, 2k+1] = c_pair[t, k, 1] = T[hi_k]
c_vals = tl.reshape(c_pair, [TILE_SIZE, BLOCK_D])
elif MSE_BITS == 4:
# 4-bit without pair LUT: exploit byte alignment
half_idx = d_offs // 2
nibble_shift = (d_offs % 2) * 4
mse_addrs = data_bases[:, None] + half_idx[None, :]
mse_raw = tl.load(
KV_cache_ptr + mse_addrs,
mask=kv_mask_1d[:, None] & d_mask[None, :],
other=0,
).to(tl.int32)
mse_idx = (mse_raw >> nibble_shift[None, :]) & 0xF
c_vals = tl.load(
Centroids_ptr + mse_idx,
mask=kv_mask_1d[:, None] & d_mask[None, :],
other=0.0,
)
else:
# Generic 3-bit path: two byte loads → 16-bit → shift/mask
mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :]
mse_raw0 = tl.load(
KV_cache_ptr + mse_addrs0,
mask=kv_mask_1d[:, None] & d_mask[None, :],
other=0,
).to(tl.int32)
mse_raw1 = tl.load(
KV_cache_ptr + mse_addrs0 + 1,
mask=kv_mask_1d[:, None] & d_mask[None, :],
other=0,
).to(tl.int32)
raw16 = mse_raw0 | (mse_raw1 << 8)
mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask_val
c_vals = tl.load(
Centroids_ptr + mse_idx,
mask=kv_mask_1d[:, None] & d_mask[None, :],
other=0.0,
)
# Opt#1: norm-correction is pre-folded into the stored per-token
# scalar at store time, so no per-tile sum+sqrt+divide here. The
# stored vec_norm carries ||k||/||c_vec||; multiply by c_vals
# below gives the original c_vec/||c_vec|| * ||k||.
# Opt#3: K-norms from per-block SoA region. Single u16 load per
# token; contiguous across tokens within a block.
norm_u16 = tl.load(
KV_cache_u16_ptr + knorm_u16_addrs, mask=kv_mask_1d, other=0
)
vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32)
# Reconstruct K: K_recon[t, d] = norm[t] * centroid[t, d]
K_recon = c_vals * vec_norms[:, None]
# K_T: [BLOCK_D, TILE_SIZE]
K_T = tl.trans(K_recon)
# S: [BLOCK_M, TILE_SIZE] via tensor core MMA
if USE_BF16_DOT:
S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16))
else:
S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16))
# Mask out-of-range positions
S = tl.where(kv_mask_1d[None, :], S, float("-inf"))
# ============================================================
# ONLINE SOFTMAX with exp2
# ============================================================
m_j = tl.maximum(M, tl.max(S, axis=1))
m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
P = tl.math.exp2(S - m_j[:, None])
l_j = tl.sum(P, axis=1)
alpha = tl.math.exp2(M - m_j)
acc = acc * alpha[:, None]
L = L * alpha + l_j
M = m_j
# ============================================================
# VALUE DEQUANT → V: [TILE_SIZE, BLOCK_D] for tl.dot(P, V)
# Opt#3: V data lives at data_base + KEY_DATA_BYTES; scale/zero
# move to the SoA metadata region (one u16 load each per token).
# ============================================================
val_bases = data_bases + KEY_DATA_BYTES
if VQB == 3:
val_addrs0 = val_bases[:, None] + val_byte_idx[None, :]
val_raw0 = tl.load(
KV_cache_ptr + val_addrs0,
mask=kv_mask_1d[:, None] & d_mask[None, :],
other=0,
).to(tl.int32)
val_raw1 = tl.load(
KV_cache_ptr + val_addrs0 + 1,
mask=kv_mask_1d[:, None] & d_mask[None, :],
other=0,
).to(tl.int32)
raw16 = val_raw0 | (val_raw1 << 8)
v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32)
else: # VQB == 4
# OPT#3 (load-halving, value path): same pattern as OPT#1 for
# keys -- load each byte exactly once as [TILE_SIZE, HALF_D],
# decode both nibbles, then tl.interleave to [TILE_SIZE, BLOCK_D]
# with pattern [v_lo0, v_hi0, v_lo1, v_hi1, ...].
V_HALF_D: tl.constexpr = BLOCK_D // 2
v_half_offs = tl.arange(0, V_HALF_D)
v_byte_mask = (v_half_offs * 2) < HEAD_DIM
val_addrs = val_bases[:, None] + v_half_offs[None, :]
val_byte = tl.load(
KV_cache_ptr + val_addrs,
mask=kv_mask_1d[:, None] & v_byte_mask[None, :],
other=0,
).to(tl.int32)
v_lo = (val_byte & 0xF).to(tl.float32) # [T, HALF_D]
v_hi = ((val_byte >> 4) & 0xF).to(tl.float32) # [T, HALF_D]
v_idx = tl.interleave(v_lo, v_hi) # [T, BLOCK_D]
scale_u16 = tl.load(
KV_cache_u16_ptr + vscale_u16_addrs, mask=kv_mask_1d, other=0
)
zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=kv_mask_1d, other=0)
v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32)
v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32)
V = v_idx * v_scales[:, None] + v_zeros[:, None]
# P·V accumulation via tensor core MMA
# P: [BLOCK_M, TILE_SIZE], V: [TILE_SIZE, BLOCK_D]
if USE_BF16_DOT:
acc += tl.dot(P.to(tl.bfloat16), V.to(tl.bfloat16))
else:
acc += tl.dot(P.to(tl.float16), V.to(tl.float16))
# ================================================================
# EPILOGUE: Store per-Q-head partial results for stage2
# ================================================================
# acc: [BLOCK_M, BLOCK_D], L: [BLOCK_M], M: [BLOCK_M]
# Stage2 expects per-Q-head: mid_o[B, Hq, split, D+1]
safe_L = tl.where(L > 0.0, L, 1.0)
one_over_L = 1.0 / safe_L
acc = acc * one_over_L[:, None]
# Convert M from log2 space back to natural log for stage2 compatibility
lse = M * LN2 + tl.log(safe_L)
# Store each Q head's result using 2D scatter
# q_head_offs: [BLOCK_M] = kv_hid * KV_GROUP_SIZE + m_offs
out_addrs = (
bid * stride_mid_b
+ q_head_offs[:, None] * stride_mid_h
+ sid * stride_mid_s
+ d_offs[None, :]
)
tl.store(
Mid_o_ptr + out_addrs,
acc,
mask=q_mask[:, None] & d_mask[None, :],
)
# Store LSE per Q head
lse_addrs = (
bid * stride_mid_b + q_head_offs * stride_mid_h + sid * stride_mid_s + HEAD_DIM
)
tl.store(Mid_o_ptr + lse_addrs, lse, mask=q_mask)
# ---------------------------------------------------------------------------
# Launcher v2
# ---------------------------------------------------------------------------
@@ -480,190 +95,3 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
shows up on a profile, cache by a hash-of-values fingerprint instead.
"""
return build_pair_lut(centroids)
def triton_turboquant_decode_attention_v2(
query: torch.Tensor,
kv_cache: torch.Tensor,
block_table: torch.Tensor,
seq_lens: torch.Tensor,
Pi: torch.Tensor,
centroids: torch.Tensor,
scale: float,
mse_bits: int,
key_packed_size: int,
value_quant_bits: int,
value_packed_size: int,
max_seq_len: int = 0, # unused; kept for backward compatibility
key_fp8: bool = False,
norm_correction: bool = False,
PiT: torch.Tensor | None = None,
# Pre-allocated buffers (optional, required for CUDA-graph stability).
mid_o_buf: torch.Tensor | None = None,
output_buf: torch.Tensor | None = None,
lse_buf: torch.Tensor | None = None,
buf_holder: Any = None,
# Fixed split count — MUST be a compile-time constant across iterations
# for CUDA-graph capture/replay to work. Mirrors the v1 launcher contract.
max_num_kv_splits: int = 32,
) -> torch.Tensor:
"""Launch optimized TQ decode attention (v2 stage1 + shared stage2).
Follows the same buffer-reuse + fixed-grid contract as the v1 launcher
(triton_turboquant_decode_attention) so the backend can capture a CUDA
graph across both versions.
"""
B, Hq, D = query.shape
Hk = kv_cache.shape[2]
block_size = kv_cache.shape[1]
padded_slot = kv_cache.shape[3]
max_num_blocks = block_table.shape[1]
n_centroids = centroids.shape[0]
kv_group_size = Hq // Hk
device = query.device
del max_seq_len # no longer used: splits is fixed via max_num_kv_splits
cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size)
# Opt#3 SoA layout constants (match store-side computation).
key_data_bytes = D if key_fp8 else cfg["mse_bytes"]
data_bytes_per_slot = key_data_bytes + cfg["val_data_bytes"]
meta_region_offset = block_size * Hk * data_bytes_per_slot
num_soa_fields = 2 if key_fp8 else 3
soa_k_norm = 0
soa_v_scale = 0 if key_fp8 else 1
soa_v_zero = 1 if key_fp8 else 2
kv_cache_u16 = kv_cache.view(torch.uint16)
# Compute q_rot = q @ Pi.T
if key_fp8:
q_rot = query.float().contiguous()
else:
q_float = query.float()
if PiT is None:
PiT = Pi.T.contiguous()
q_rot = (q_float @ PiT).contiguous()
# BLOCK_M: pad KV_GROUP_SIZE to power of 2, minimum 16 for tensor cores
BLOCK_M = max(16, triton.next_power_of_2(kv_group_size))
# TILE_SIZE (BLOCK_KV): tokens per inner-loop iteration
TILE_SIZE = 16
# Fixed split count — must be constant across calls for cudagraph replay.
NUM_KV_SPLITS = max_num_kv_splits
# --- mid_o buffer reuse (same pattern as v1 launcher) ---
if (
mid_o_buf is not None
and mid_o_buf.shape[0] >= B
and mid_o_buf.shape[2] >= NUM_KV_SPLITS
):
mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :]
else:
mid_o = torch.empty(
B,
Hq,
NUM_KV_SPLITS,
D + 1,
dtype=torch.float32,
device=device,
)
if buf_holder is not None:
buf_holder._tq_mid_o_buf = mid_o
fp8_e4b15 = _use_fp8_e4b15(device.index or 0)
# Build pair LUT for 4-bit MSE (FLUTE §3.2)
use_pair_lut = mse_bits == 4 and not key_fp8
pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids
# Platform-dependent pipelining depth: ROCm prefers num_stages=1.
stage1_num_stages = 1 if _is_hip else 2
# --- Stage 1: v2 kernel ---
# Grid over KV heads (not Q heads) — each program handles all Q heads
grid = (B, Hk, NUM_KV_SPLITS)
_tq_decode_stage1_v2[grid](
q_rot,
kv_cache,
kv_cache_u16,
block_table,
seq_lens,
centroids,
pair_lut,
mid_o,
q_rot.stride(0),
q_rot.stride(1),
kv_cache.stride(0),
block_table.stride(0),
mid_o.stride(0),
mid_o.stride(1),
mid_o.stride(2),
NUM_Q_HEADS=Hq,
NUM_KV_HEADS=Hk,
HEAD_DIM=D,
BLOCK_SIZE=block_size,
PADDED_SLOT=padded_slot,
MAX_NUM_BLOCKS=max_num_blocks,
NUM_KV_SPLITS=NUM_KV_SPLITS,
KV_GROUP_SIZE=kv_group_size,
MSE_BITS=mse_bits,
MSE_BYTES=cfg["mse_bytes"],
VQB=value_quant_bits,
VAL_DATA_BYTES=cfg["val_data_bytes"],
KEY_DATA_BYTES=key_data_bytes,
META_REGION_OFFSET=meta_region_offset,
NUM_SOA_FIELDS=num_soa_fields,
SOA_K_NORM=soa_k_norm,
SOA_V_SCALE=soa_v_scale,
SOA_V_ZERO=soa_v_zero,
N_CENTROIDS=n_centroids,
ATTN_SCALE=scale,
BLOCK_D=cfg["BLOCK_D"],
TILE_SIZE=TILE_SIZE,
BLOCK_M=BLOCK_M,
KEY_FP8=1 if key_fp8 else 0,
NORM_CORRECTION=1 if norm_correction else 0,
FP8_E4B15=fp8_e4b15,
USE_PAIR_LUT=1 if use_pair_lut else 0,
USE_BF16_DOT=1 if _is_hip else 0,
num_warps=4,
num_stages=stage1_num_stages,
)
# --- output / lse buffer reuse (same pattern as v1 launcher) ---
if output_buf is not None and output_buf.shape[0] >= B:
output = output_buf[:B, :Hq, :D]
else:
output = torch.empty(B, Hq, D, dtype=torch.float32, device=device)
if buf_holder is not None:
buf_holder._tq_output_buf = output
if lse_buf is not None and lse_buf.shape[0] >= B:
lse = lse_buf[:B, :Hq]
else:
lse = torch.empty(B, Hq, dtype=torch.float32, device=device)
if buf_holder is not None:
buf_holder._tq_lse_buf = lse
# --- Stage 2: reduce across KV splits (unchanged) ---
grid2 = (B, Hq)
_fwd_kernel_stage2[grid2](
mid_o,
output,
lse,
seq_lens,
mid_o.stride(0),
mid_o.stride(1),
mid_o.stride(2),
output.stride(0),
output.stride(1),
lse.stride(0),
NUM_KV_SPLITS=NUM_KV_SPLITS,
BLOCK_DV=cfg["BLOCK_D"],
Lv=D,
num_warps=4,
num_stages=2,
)
return output.to(query.dtype)
vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py
-15
View File
@@ -876,21 +876,6 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
return build_pair_lut(centroids)
# PiT in Q's dtype cached per (PiT-id, dtype) to avoid a host cast on every
# launcher call. Keyed by the fp32 PiT's data_ptr so ownership/lifetime is
# tied to the caller's tensor.
_pit_qdtype_cache: dict = {}
def _get_pit_in_query_dtype(PiT: torch.Tensor, qdtype: torch.dtype) -> torch.Tensor:
key = (PiT.data_ptr(), qdtype)
cached = _pit_qdtype_cache.get(key)
if cached is None:
cached = PiT.to(qdtype).contiguous()
_pit_qdtype_cache[key] = cached
return cached
def triton_turboquant_unified_attention(
query: torch.Tensor, # [num_tokens, Hq, D] - fp16/bf16
kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8