diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py index aff4fc1c3f..3cee44ba50 100644 --- a/vllm/model_executor/layers/mamba/ops/replayssm_config.py +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -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))) diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index 64e6fd46d7..09e7a40aa5 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -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. diff --git a/vllm/models/inkling/amd/ops/lamport.py b/vllm/models/inkling/amd/ops/lamport.py deleted file mode 100644 index 9705a9a3df..0000000000 --- a/vllm/models/inkling/amd/ops/lamport.py +++ /dev/null @@ -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 diff --git a/vllm/models/inkling/amd/ops/mm_towers.py b/vllm/models/inkling/amd/ops/mm_towers.py deleted file mode 100644 index 021cf5274b..0000000000 --- a/vllm/models/inkling/amd/ops/mm_towers.py +++ /dev/null @@ -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) diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py index 82782fcb96..7baec0b42b 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py @@ -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] diff --git a/vllm/v1/attention/ops/triton_fp8_mqa_logits.py b/vllm/v1/attention/ops/triton_fp8_mqa_logits.py deleted file mode 100644 index 619d0ec50a..0000000000 --- a/vllm/v1/attention/ops/triton_fp8_mqa_logits.py +++ /dev/null @@ -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 diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py index 321f42bcb1..68ce0e0de2 100644 --- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py +++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py @@ -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) diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py index 03ef36e5b6..d77592010a 100644 --- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py +++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py @@ -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