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

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

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

Synced Nov 4, 06:15 UTC

[New Model][ROCm] Add AMD support for DeepSeek V4 (#40871)

Signed-off-by: ganyi <ygan@amd.com> Signed-off-by: whx-sjtu <xiaowang990929@gmail.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Signed-off-by: tjtanaavllm <tunjian.tan@amd.com> Co-authored-by: ganyi <ygan@amd.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaavllm <tunjian.tan@amd.com>

Hexiang Wang committed May 5, 2026 at 23:55 UTC 628c43630155b9fa8a8ec8871eb90c06bf8b36e6
22 files changed +939 -134
CMakeLists.txt
+6 -6
@@ -307,12 +307,12 @@ set(VLLM_EXT_SRC
307 "csrc/quantization/activation_kernels.cu"
308 "csrc/cuda_utils_kernels.cu"
309 "csrc/custom_all_reduce.cu"
310 - "csrc/torch_bindings.cpp")
310 + "csrc/torch_bindings.cpp"
311 + "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
312
313 if(VLLM_GPU_LANG STREQUAL "CUDA")
314 list(APPEND VLLM_EXT_SRC
314 - "csrc/minimax_reduce_rms_kernel.cu"
315 - "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
315 + "csrc/minimax_reduce_rms_kernel.cu")
316
317 SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
318
@@ -1047,13 +1047,13 @@ endif()
1047 set(VLLM_MOE_EXT_SRC
1048 "csrc/moe/torch_bindings.cpp"
1049 "csrc/moe/moe_align_sum_kernels.cu"
1050 - "csrc/moe/topk_softmax_kernels.cu")
1050 + "csrc/moe/topk_softmax_kernels.cu"
1051 + "csrc/moe/topk_softplus_sqrt_kernels.cu")
1052
1053 if(VLLM_GPU_LANG STREQUAL "CUDA")
1054 list(APPEND VLLM_MOE_EXT_SRC
1055 "csrc/moe/moe_wna16.cu"
1055 - "csrc/moe/grouped_topk_kernels.cu"
1056 - "csrc/moe/topk_softplus_sqrt_kernels.cu")
1056 + "csrc/moe/grouped_topk_kernels.cu")
1057 endif()
1058
1059 if(VLLM_GPU_LANG STREQUAL "CUDA")
csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+36 -2
@@ -29,7 +29,11 @@
29 */
30
31 #include <cmath>
32 -#include <cuda_fp8.h>
32 +#ifndef USE_ROCM
33 + #include <cuda_fp8.h>
34 +#else
35 + #include <hip/hip_fp8.h>
36 +#endif
37 #include <cuda_runtime.h>
38 #include <type_traits>
39
@@ -42,7 +46,23 @@
46 #include "type_convert.cuh"
47
48 #ifndef FINAL_MASK
45 - #define FINAL_MASK 0xffffffffu
49 + #ifdef USE_ROCM
50 + #define FINAL_MASK 0xffffffffffffffffULL
51 + #else
52 + #define FINAL_MASK 0xffffffffu
53 + #endif
54 +#endif
55 +
56 +#ifdef USE_ROCM
57 +// ROCm-compatible FP8 conversion helpers
58 +__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
59 + #if defined(HIP_FP8_TYPE_OCP)
60 + __hip_fp8_e4m3 fp8_val(val);
61 + #else
62 + __hip_fp8_e4m3_fnuz fp8_val(val);
63 + #endif
64 + return reinterpret_cast<uint8_t&>(fp8_val);
65 +}
66 #endif
67
68 namespace vllm {
@@ -314,9 +334,13 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
334 for (int i = 0; i < kElemsPerLane; i++) {
335 float scaled = elements[i] * inv_scale;
336 scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
337 +#ifndef USE_ROCM
338 __nv_fp8_storage_t s =
339 __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
340 out_bytes[i] = static_cast<uint8_t>(s);
341 +#else
342 + out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
343 +#endif
344 }
345 // One 16-byte STG per lane.
346 *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
@@ -384,6 +408,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
408 // PDL: enable programmatic stream serialization whenever the hardware
409 // supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable,
410 // so leave numAttrs = 0 and launch as a regular kernel.
411 +#ifndef USE_ROCM
412 static int const sm_version = getSMVersion();
413 // Host-side guard: the device kernel body is compiled as a no-op for
414 // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
@@ -410,6 +435,15 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
435 q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
436 num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
437 kv_block_stride);
438 +#else
439 + // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
440 + // clang-format off
441 + fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
442 + <<<grid, kBlockSize, 0, stream>>>(
443 + q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
444 + eps, num_tokens_full, num_tokens_insert, num_heads_q,
445 + cache_block_size, kv_block_stride);
446 +#endif
447 }
448
449 } // namespace deepseek_v4_fused_ops
csrc/moe/topk_softplus_sqrt_kernels.cu
+32 -21
@@ -60,15 +60,6 @@ __device__ __forceinline__ float toFloat(T value) {
60 }
61 }
62
63 -#define FINAL_MASK 0xffffffff
64 -template <typename T>
65 -__inline__ __device__ T warpReduceSum(T val) {
66 -#pragma unroll
67 - for (int mask = 16; mask > 0; mask >>= 1)
68 - val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
69 - return val;
70 -}
71 -
63 // ====================== TopK softplus_sqrt things
64 // ===============================
65
@@ -272,8 +263,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
263 }
264 }
265 // Compute per-thread scale (using warp reduction when renormalizing).
266 + // THREADS_PER_ROW-parameterized butterfly works for both warp sizes (32
267 + // on CUDA, 64 on ROCm CDNA) and any THREADS_PER_ROW the dispatch picks.
268 if (renormalize) {
276 - selected_sum = warpReduceSum(selected_sum);
269 +#pragma unroll
270 + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
271 + selected_sum +=
272 + VLLM_SHFL_XOR_SYNC_WIDTH(selected_sum, mask, THREADS_PER_ROW);
273 + }
274 }
275 float scale = static_cast<float>(routed_scaling_factor);
276 if (renormalize) {
@@ -544,7 +541,6 @@ void topkGatingSoftplusSqrtKernelLauncher(
541 const IndType* tid2eid, cudaStream_t stream) {
542 static constexpr int WARPS_PER_TB = 4;
543 static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
547 -#ifndef USE_ROCM
544 // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
545 // elements can be loaded by a warp
546 static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
@@ -552,6 +548,19 @@ void topkGatingSoftplusSqrtKernelLauncher(
548 std::is_same_v<InputType, __half>)
549 ? 4
550 : 8;
551 + // Narrower LDG (ELTS_PER_LDG=1) used by 192/320/448/576 on ROCm WARP_SIZE=64
552 + // where ELTS_PER_LDG=2 fails the EXPERTS%(ELTS_PER_LDG*WARP_SIZE)==0 check.
553 + // On CUDA WARP_SIZE=32 the wider LDG already aligns, so the alias collapses
554 + // back to BYTES_PER_LDG_MULTIPLE_64 — no behavioral change for CUDA.
555 +#ifdef USE_ROCM
556 + static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW =
557 + (std::is_same_v<InputType, __nv_bfloat16> ||
558 + std::is_same_v<InputType, __half>)
559 + ? 2
560 + : 4;
561 +#else
562 + static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW =
563 + BYTES_PER_LDG_MULTIPLE_64;
564 #endif
565 switch (num_experts) {
566 case 1:
@@ -584,27 +593,29 @@ void topkGatingSoftplusSqrtKernelLauncher(
593 case 512:
594 LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
595 break;
587 - // (CUDA only) support multiples of 64 when num_experts is not power of 2.
588 - // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
589 - // num_experts, alternatively we can test 4 bytes loading and enable it in
590 - // future.
591 -#ifndef USE_ROCM
596 + // Multiples of 64 that are not powers of 2. The kernel requires
597 + // EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0. With ELTS_PER_LDG=2
598 + // (BYTES_PER_LDG_MULTIPLE_64), this holds for all five values on CUDA
599 + // WARP_SIZE=32 but only for 384 on ROCm WARP_SIZE=64. The other four
600 + // use BYTES_PER_LDG_MULTIPLE_64_NARROW (ELTS_PER_LDG=1), which
601 + // satisfies the assertion for any multiple of 64 on either backend;
602 + // on CUDA the narrow alias collapses back to the wider load, so CUDA
603 + // behavior is unchanged.
604 case 192:
593 - LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
605 + LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
606 break;
607 case 320:
596 - LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
608 + LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
609 break;
610 case 384:
611 LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
612 break;
613 case 448:
602 - LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
614 + LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
615 break;
616 case 576:
605 - LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
617 + LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
618 break;
607 -#endif
619 default: {
620 TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
621 }
csrc/moe/torch_bindings.cpp
+1 -2
@@ -16,14 +16,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
16 "bias) -> ()");
17 m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
18
19 -#ifndef USE_ROCM
19 m.def(
20 "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
21 "token_expert_indices, Tensor gating_output, bool renormalize, float "
22 "routed_scaling_factor, Tensor? "
23 "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
24 m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
26 -#endif
25 +
26 // Calculate the result of moe by summing up the partial results
27 // from all selected experts.
28 m.def("moe_sum(Tensor input, Tensor! output) -> ()");
csrc/torch_bindings.cpp
-2
@@ -183,7 +183,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
183 "int forced_token_heads_per_warp=-1) -> ()");
184 ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
185
186 -#ifndef USE_ROCM
186 // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
187 // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
188 // kernel launch.
@@ -194,7 +193,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
193 "float eps, int cache_block_size) -> ()");
194 ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
195 &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
197 -#endif
196
197 // Apply repetition penalties to logits in-place
198 ops.def(
requirements/rocm.txt
+3
@@ -21,3 +21,6 @@ timm>=1.0.17
21 # amd-quark: required for Quark quantization on ROCm
22 # To be consistent with test_quark.py
23 amd-quark>=0.8.99
24 +# tilelang has to be installed for mhc module to be
25 +# imported correctly.
26 +tilelang==0.1.9
tests/kernels/moe/test_topk_softplus_sqrt.py
+4 -2
@@ -70,7 +70,8 @@ def test_sqrtsoftplus_bias_uses_deepseek_v4_routing_method():
70
71
72 @pytest.mark.skipif(
73 - not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
73 + not current_platform.is_cuda_alike(),
74 + reason="This test is skipped on non-CUDA platform.",
75 )
76 @pytest.mark.parametrize("num_tokens", [1, 33, 128])
77 @pytest.mark.parametrize("hidden_size", [1024, 2048])
@@ -125,7 +126,8 @@ def test_fused_topk_softplus_sqrt(
126
127
128 @pytest.mark.skipif(
128 - not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
129 + not current_platform.is_cuda_alike(),
130 + reason="This test is skipped on non-CUDA platform.",
131 )
132 @pytest.mark.parametrize("num_tokens", [1, 33, 128])
133 @pytest.mark.parametrize("hidden_size", [1024, 2048])
vllm/config/kernel.py
+2
@@ -119,6 +119,7 @@ MoEBackend = Literal[
119 "flashinfer_cutedsl",
120 "marlin",
121 "humming",
122 + "triton_unfused",
123 "aiter",
124 "emulation",
125 ]
@@ -150,6 +151,7 @@ class KernelConfig:
151 - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)
152 - "marlin": Use Marlin kernels (weight-only quantization)
153 - "humming": Use Humming Mixed Precision kernels
154 + - "triton_unfused": Use Triton unfused MoE kernels
155 - "aiter": Use AMD AITer kernels (ROCm only)
156 - "emulation": use BF16/FP16 GEMM, dequantizing weights and
157 running QDQ on activations.
vllm/model_executor/kernels/linear/scaled_mm/aiter.py
+15
@@ -312,6 +312,21 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
312 As: torch.Tensor,
313 Bs: torch.Tensor,
314 ) -> torch.Tensor:
315 + if As.dtype != Bs.dtype:
316 + from vllm.model_executor.layers.quantization.utils.fp8_utils import (
317 + _upcast_e8m0_to_fp32,
318 + )
319 +
320 + if As.dtype == torch.float8_e8m0fnu:
321 + As = _upcast_e8m0_to_fp32(As).contiguous()
322 + else:
323 + As = As.to(torch.float32)
324 +
325 + if Bs.dtype == torch.float8_e8m0fnu:
326 + Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
327 + else:
328 + Bs = Bs.to(torch.float32)
329 +
330 out_dtype = self.config.out_dtype
331 if self.use_triton:
332 gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale
vllm/model_executor/layers/activation.py
+3 -1
@@ -169,7 +169,9 @@ class SiluAndMulWithClamp(CustomOp):
169 def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
170 super().__init__(compile_native=compile_native)
171 self.swiglu_limit = float(swiglu_limit)
172 - if current_platform.is_cuda_alike() or current_platform.is_xpu():
172 + if current_platform.is_rocm():
173 + self._forward_method = self.forward_native
174 + elif current_platform.is_cuda_alike() or current_platform.is_xpu():
175 self.op = torch.ops._C.silu_and_mul_with_clamp
176 elif current_platform.is_cpu():
177 self._forward_method = self.forward_native
vllm/model_executor/layers/deepseek_compressor.py
+3 -2
@@ -300,6 +300,7 @@ class DeepseekCompressor(nn.Module):
300 state_cache = self.state_cache.kv_cache
301 # kv_state stored in first half, score_state stored in second half
302 state_width = state_cache.shape[-1] // 2
303 + pdl_kwargs = {} if current_platform.is_rocm() else {"launch_pdl": False}
304
305 # Store the KV and score (with fused APE addition) in the state.
306 # NOTE: PDL is disabled — both this kernel and _fused_kernel below
@@ -324,7 +325,7 @@ class DeepseekCompressor(nn.Module):
325 TRITON_BLOCK_SIZE=triton.next_power_of_2(kv.shape[-1]),
326 STATE_WIDTH=state_width,
327 COMPRESS_RATIO=self.compress_ratio,
327 - launch_pdl=False,
328 + **pdl_kwargs,
329 )
330
331 # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write.
@@ -373,7 +374,7 @@ class DeepseekCompressor(nn.Module):
374 SCALE_DIM=self._scale_dim,
375 KV_BLOCK_STRIDE=kv_cache.stride(0),
376 num_warps=self._num_warps,
376 - launch_pdl=False,
377 + **pdl_kwargs,
378 )
379
380
vllm/model_executor/layers/deepseek_v4_attention.py
+73 -19
@@ -28,6 +28,11 @@ from vllm.v1.attention.ops.deepseek_v4_ops import (
28 fused_inv_rope_fp8_quant,
29 fused_q_kv_rmsnorm,
30 )
31 +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
32 + rocm_forward_decode_fallback,
33 + rocm_inv_rope_einsum,
34 + rocm_sparse_attn_prefill,
35 +)
36
37 if TYPE_CHECKING:
38 from vllm.v1.attention.backends.mla.sparse_swa import (
@@ -53,6 +58,7 @@ from vllm.model_executor.layers.quantization.input_quant_fp8 import (
58 from vllm.model_executor.layers.quantization.utils.quant_utils import (
59 GroupShape,
60 )
61 +from vllm.platforms import current_platform
62 from vllm.utils.multi_stream_utils import (
63 execute_in_parallel,
64 maybe_execute_in_parallel,
@@ -198,8 +204,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
204 # Pick fp8_einsum recipe based on GPU arch:
205 # SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128
206 # SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1
201 - from vllm.platforms import current_platform
202 -
207 cap = current_platform.get_device_capability()
208 assert cap is not None, "DeepseekV4 attention requires a CUDA device"
209 self._einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128)
@@ -222,6 +226,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
226 + 1 # 1B pad
227 )
228
229 + # Will be None on ROCm for now.
230 self.aux_stream_list = mla_modules.aux_stream_list
231 # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
232 # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins
@@ -303,6 +308,19 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
308 )
309 o = o_padded[:, : self.n_local_heads, :]
310
311 + # Keep ROCm on the BF16 reference wo_a path util kernel ready.
312 + if current_platform.is_rocm():
313 + z = rocm_inv_rope_einsum(
314 + self.rotary_emb,
315 + o,
316 + positions,
317 + self.rope_head_dim,
318 + self.n_local_groups,
319 + self.o_lora_rank,
320 + self.wo_a,
321 + )
322 + return self.wo_b(z.flatten(1))
323 +
324 # O projection: inverse RoPE + FP8 quant + einsum + wo_b
325 o_fp8, o_scale = fused_inv_rope_fp8_quant(
326 o,
@@ -336,12 +354,15 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
354 return self.wo_b(z.flatten(1))
355
356 def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]:
339 - assert self.aux_stream_list is not None
340 - assert len(self.aux_stream_list) >= 3
357 + aux_streams = self.aux_stream_list
358 + if aux_streams is not None:
359 + assert len(aux_streams) >= 3
360 + aux_streams = aux_streams[:3]
361
362 # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs
363 # on aux streams 0..2 when their owning module exists. ln_events[0]
364 # is the fan-out start event; ln_events[1..3] are per-aux done events.
365 + # On ROCm, aux_streams is None and execute_in_parallel runs serially.
366 aux_fns: list[Callable[[], Any] | None] = [None, None, None]
367
368 if self.compressor is not None:
@@ -385,7 +406,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
406 aux_fns,
407 self.ln_events[0],
408 self.ln_events[1:4],
388 - self.aux_stream_list[:3],
409 + aux_streams,
410 enable=hidden_states.shape[0]
411 <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD,
412 )
@@ -419,8 +440,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
440 # downstream reads q on default). Indexer/compressor go on aux for
441 # overlap with default's GEMM + cache write.
442 if self.indexer is not None:
422 - assert self.aux_stream_list is not None
423 - aux_stream = self.aux_stream_list[0]
443 + aux_stream = (
444 + self.aux_stream_list[0] if self.aux_stream_list is not None else None
445 + )
446 indexer = self.indexer
447 # Local ref so the closure keeps a non-None type for mypy.
448 assert self.compressor is not None
@@ -448,8 +470,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
470 )
471 elif self.compressor is not None:
472 # wq_b + kv_insert on default, compressor on aux.
451 - assert self.aux_stream_list is not None
452 - aux_stream = self.aux_stream_list[0]
473 + aux_stream = (
474 + self.aux_stream_list[0] if self.aux_stream_list is not None else None
475 + )
476 compressor = self.compressor
477
478 def wq_b_kv_insert() -> torch.Tensor:
@@ -668,7 +691,7 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
691 vllm_config.scheduler_config.max_num_batched_tokens
692 )
693 self.max_model_len = vllm_config.model_config.max_model_len
671 - # DeepseekV4 only supports fp8 kv-cache format for now
694 + # DeepseekV4 only supports fp8 kv-cache format for now.
695 kv_cache_dtype = cache_config.cache_dtype if cache_config is not None else "fp8"
696
697 assert kv_cache_dtype.startswith("fp8"), (
@@ -816,6 +839,25 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
839 swa_indices = swa_metadata.decode_swa_indices
840 swa_lens = swa_metadata.decode_swa_lens
841
842 + if current_platform.is_rocm():
843 + rocm_forward_decode_fallback(
844 + q=q,
845 + kv_cache=kv_cache,
846 + swa_k_cache=self.swa_cache_layer.kv_cache,
847 + swa_only=swa_only,
848 + topk_indices=topk_indices,
849 + topk_lens=topk_lens,
850 + swa_indices=swa_indices,
851 + swa_lens=swa_lens,
852 + attn_sink=self.attn_sink,
853 + scale=self.scale,
854 + head_dim=self.head_dim,
855 + nope_head_dim=self.nope_head_dim,
856 + rope_head_dim=self.rope_head_dim,
857 + output=output,
858 + )
859 + return
860 +
861 # We treat queries in the same seq as different queries
862 # and later we only attend by generated indices.
863 # q arrives pre-padded to self.padded_heads by the outer wrapper.
@@ -980,15 +1022,27 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
1022 N,
1023 )
1024
983 - output_chunk, _, _ = flash_mla_sparse_fwd(
984 - q=q[query_start:query_end],
985 - kv=kv.view(-1, 1, q.shape[-1]),
986 - indices=combined_indices.unsqueeze(1),
987 - sm_scale=self.scale,
988 - attn_sink=self.attn_sink,
989 - topk_length=combined_lens,
990 - out=output[query_start:query_end],
991 - )
1025 + if current_platform.is_rocm():
1026 + rocm_sparse_attn_prefill(
1027 + q=q[query_start:query_end],
1028 + kv=kv.view(-1, 1, q.shape[-1]),
1029 + indices=combined_indices.unsqueeze(1),
1030 + topk_length=combined_lens,
1031 + scale=self.scale,
1032 + head_dim=self.head_dim,
1033 + attn_sink=self.attn_sink,
1034 + output=output[query_start:query_end],
1035 + )
1036 + else:
1037 + output_chunk, _, _ = flash_mla_sparse_fwd(
1038 + q=q[query_start:query_end],
1039 + kv=kv.view(-1, 1, q.shape[-1]),
1040 + indices=combined_indices.unsqueeze(1),
1041 + sm_scale=self.scale,
1042 + attn_sink=self.attn_sink,
1043 + topk_length=combined_lens,
1044 + out=output[query_start:query_end],
1045 + )
1046
1047
1048 class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+79 -2
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
18 from vllm.model_executor.layers.fused_moe.config import (
19 FusedMoEQuantConfig,
20 FusedMoEQuantDesc,
21 + RoutingMethodType,
22 mxfp4_mxfp8_moe_quant_config,
23 mxfp4_w4a8_moe_quant_config,
24 mxfp4_w4a16_moe_quant_config,
@@ -64,6 +65,8 @@ class Mxfp4MoeBackend(Enum):
65 MARLIN = "MARLIN"
66 # ROCm AITER backends
67 AITER_MXFP4_BF16 = "AITER_MXFP4_BF16" # W4A16: CK kernel
68 + # Keep the legacy name as an alias while the ROCm split backend rename settles.
69 + AITER = "AITER_MXFP4_BF16"
70 AITER_MXFP4_FP8 = "AITER_MXFP4_FP8" # W4A8: triton kernel
71 # Triton
72 TRITON = "TRITON"
@@ -253,6 +256,8 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
256 TRTLLM MXFP8; SM90 falls through to Triton_unfused or Marlin (the
257 backend-level ``is_supported_config`` check filters by device capability).
258 """
259 + if current_platform.is_rocm():
260 + return [Mxfp4MoeBackend.AITER_MXFP4_BF16]
261 _AVAILABLE_BACKENDS = [
262 Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
263 Mxfp4MoeBackend.DEEPGEMM_MXFP4,
@@ -543,8 +548,22 @@ def select_deepseek_v4_mxfp4_moe_backend(
548 activation_format,
549 )
550
551 + # DeepSeek-V4 on ROCm is more accurate with the unfused Triton MXFP4 path
552 + # than the default AITER path. Prefer Triton-unfused for this routing mode,
553 + # while keeping AITER as a fallback if Triton-unfused rejects the config.
554 + if (
555 + current_platform.is_rocm()
556 + and config.routing_method == RoutingMethodType.DeepseekV4
557 + ):
558 + priority_backends = [
559 + Mxfp4MoeBackend.TRITON_UNFUSED,
560 + Mxfp4MoeBackend.AITER_MXFP4_BF16,
561 + ]
562 + else:
563 + priority_backends = _get_priority_backends()
564 +
565 # Iterate priority backends: TRTLLM MXFP8, then Triton.
547 - for backend in _get_priority_backends():
566 + for backend in priority_backends:
567 activation_key = _backend_activation_key(backend)
568 for k_cls in backend_to_kernel_cls(backend):
569 supported, reason = k_cls.is_supported_config(
@@ -1252,6 +1271,64 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1271 w2_bias,
1272 )
1273
1274 + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
1275 + from vllm._aiter_ops import rocm_aiter_ops
1276 +
1277 + if w13_bias is not None:
1278 + w13_bias = w13_bias.data.to(torch.float32)
1279 + if w2_bias is not None:
1280 + w2_bias = w2_bias.data.to(torch.float32)
1281 +
1282 + e, n, k = w13_weight.shape
1283 +
1284 + w13_weight.view(torch.uint8).copy_(
1285 + w13_weight.data.view(torch.uint8)
1286 + .view(e, n // 2, 2, k)
1287 + .permute(0, 2, 1, 3)
1288 + .contiguous()
1289 + .view(e, n, k)
1290 + )
1291 + w13_weight_scale.data = (
1292 + w13_weight_scale.data.view(e, n // 2, 2, -1)
1293 + .permute(0, 2, 1, 3)
1294 + .contiguous()
1295 + .view(e, n, -1)
1296 + )
1297 +
1298 + w13_weight.data = w13_weight.data.view(torch.float4_e2m1fn_x2)
1299 + w2_weight.data = w2_weight.data.view(torch.float4_e2m1fn_x2)
1300 +
1301 + w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w13_weight, 16, True)
1302 + shuffled_w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(
1303 + w13_weight_scale.view(-1, w13_weight_scale.shape[-1]),
1304 + num_experts,
1305 + True,
1306 + )
1307 +
1308 + w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w2_weight, 16, False)
1309 + shuffled_w2_scale = rocm_aiter_ops.shuffle_scale_a16w4(
1310 + w2_weight_scale.view(-1, w2_weight_scale.shape[-1]),
1311 + num_experts,
1312 + False,
1313 + )
1314 +
1315 + if w13_bias is not None:
1316 + w13_bias = (
1317 + w13_bias.data.view(-1, n // 2, 2)
1318 + .permute(0, 2, 1)
1319 + .contiguous()
1320 + .view(-1, n)
1321 + )
1322 +
1323 + return (
1324 + w13_weight,
1325 + w2_weight,
1326 + shuffled_w13_scale,
1327 + shuffled_w2_scale,
1328 + w13_bias,
1329 + w2_bias,
1330 + )
1331 +
1332 elif mxfp4_backend in TRITON_BACKENDS:
1333 from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1334
@@ -1307,7 +1384,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1384 else:
1385 raise ValueError(
1386 f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
1310 - f"Expected TRTLLM or Triton backend."
1387 + f"Expected TRTLLM, Triton, or AITER backend."
1388 )
1389
1390
vllm/model_executor/layers/mhc.py
+105 -2
@@ -234,6 +234,39 @@ def mhc_pre(
234 num_tokens = residual_flat.shape[0]
235 fn_flat = fn
236
237 + if current_platform.is_rocm():
238 + x = residual_flat.view(num_tokens, hc_mult * hidden_size).to(torch.float32)
239 + mixes = torch.matmul(x, fn_flat.t())
240 + sqrsum = x.square().sum(dim=-1, keepdim=True)
241 + mixes = mixes * torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps)
242 +
243 + pre_logits = mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult]
244 + pre_mix = torch.sigmoid(pre_logits) + hc_pre_eps
245 +
246 + post_logits = (
247 + mixes[:, hc_mult : 2 * hc_mult] * hc_scale[1]
248 + + hc_base[hc_mult : 2 * hc_mult]
249 + )
250 + post_mix = torch.sigmoid(post_logits) * hc_post_mult_value
251 +
252 + comb_logits = mixes[:, 2 * hc_mult :].view(
253 + num_tokens, hc_mult, hc_mult
254 + ) * hc_scale[2] + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult)
255 + comb_mix = torch.softmax(comb_logits, dim=-1) + hc_sinkhorn_eps
256 + comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
257 + for _ in range(sinkhorn_repeat - 1):
258 + comb_mix = comb_mix / (comb_mix.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps)
259 + comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
260 +
261 + layer_input = torch.sum(
262 + pre_mix.unsqueeze(-1) * residual_flat.to(torch.float32), dim=1
263 + ).to(torch.bfloat16)
264 + return (
265 + post_mix.view(*outer_shape, hc_mult, 1),
266 + comb_mix.view(*outer_shape, hc_mult, hc_mult),
267 + layer_input.view(*outer_shape, hidden_size),
268 + )
269 +
270 # these number are from deepgemm kernel impl
271 block_k = 64
272 block_m = 64
@@ -414,6 +447,14 @@ def mhc_post(
447 post_layer_mix: torch.Tensor,
448 comb_res_mix: torch.Tensor,
449 ) -> torch.Tensor:
450 + if current_platform.is_rocm():
451 + mixed_residual = torch.einsum(
452 + "...ij,...ih->...jh",
453 + comb_res_mix.to(torch.float32),
454 + residual.to(torch.float32),
455 + )
456 + post_term = post_layer_mix.to(torch.float32) * x.unsqueeze(-2).to(torch.float32)
457 + return (mixed_residual + post_term).to(residual.dtype)
458 out = torch.empty_like(residual)
459 mhc_post_tilelang(
460 comb_res_mix,
@@ -551,6 +592,49 @@ def hc_head_fuse_tilelang(
592 T.pdl_trigger()
593
594
595 +def _hc_head_fused_reference(
596 + hs_flat: torch.Tensor,
597 + fn: torch.Tensor,
598 + hc_scale: torch.Tensor,
599 + hc_base: torch.Tensor,
600 + out: torch.Tensor,
601 + hidden_size: int,
602 + rms_eps: float,
603 + hc_eps: float,
604 + hc_mult: int,
605 +) -> None:
606 + """Pure-PyTorch reference for `hc_head_fuse_tilelang`.
607 +
608 + Used on platforms where the tilelang HIP/CUDA backend is not available
609 + (e.g. ROCm builds shipping a tilelang wheel without `target.build.tilelang_hip`).
610 + Mirrors the math of the tilelang kernel exactly:
611 +
612 + x = hs_flat.flatten(-2, -1) # (T, hc_mult * H), fp32
613 + mixes = x @ fn.T # (T, hc_mult)
614 + rsqrt = 1 / sqrt(||x||^2 / (hc_mult * H) + rms_eps)
615 + pre[m] = sigmoid(mixes[m] * rsqrt * hc_scale[0] + hc_base[m]) + hc_eps
616 + out = sum_m pre[m] * hs_flat[:, m, :] # cast back to bf16
617 +
618 + `out` is mutated in place to keep the same op contract
619 + (`mutates_args=["out"]`).
620 + """
621 + num_tokens = hs_flat.shape[0]
622 + if num_tokens == 0:
623 + return
624 + x = hs_flat.reshape(num_tokens, hc_mult * hidden_size).to(torch.float32)
625 + # fn: (hc_mult, hc_mult * hidden_size) → mixes: (T, hc_mult)
626 + mixes = torch.matmul(x, fn.t())
627 + sqrsum = x.square().sum(dim=-1, keepdim=True)
628 + rsqrt = torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps)
629 + # hc_scale has shape (1,); hc_base has shape (hc_mult,)
630 + pre_mix = torch.sigmoid(mixes * rsqrt * hc_scale[0] + hc_base) + hc_eps
631 + # weighted sum over the hc_mult channel dim
632 + result = torch.sum(pre_mix.unsqueeze(-1) * hs_flat.to(torch.float32), dim=1).to(
633 + out.dtype
634 + )
635 + out.copy_(result)
636 +
637 +
638 def _hc_head_fused_kernel(
639 hs_flat: torch.Tensor,
640 fn: torch.Tensor,
@@ -563,8 +647,15 @@ def _hc_head_fused_kernel(
647 hc_mult: int,
648 ) -> None:
649 """Fill pre-allocated `out` (T, H) in-place with the hc_head result."""
566 - if hs_flat.shape[0] > 0:
567 - hc_head_fuse_tilelang(
650 + if hs_flat.shape[0] == 0:
651 + return
652 + if current_platform.is_rocm():
653 + # tilelang ships only the CUDA codegen in upstream wheels, so the HIP
654 + # FFI target (`target.build.tilelang_hip`) is missing and the JIT call
655 + # would raise `ValueError: Cannot find global function ...`. Use a
656 + # numerically equivalent torch fallback instead. `mhc_pre` and
657 + # `mhc_post` already follow this same pattern above.
658 + _hc_head_fused_reference(
659 hs_flat,
660 fn,
661 hc_scale,
@@ -575,6 +666,18 @@ def _hc_head_fused_kernel(
666 hc_eps,
667 hc_mult,
668 )
669 + return
670 + hc_head_fuse_tilelang(
671 + hs_flat,
672 + fn,
673 + hc_scale,
674 + hc_base,
675 + out,
676 + hidden_size,
677 + rms_eps,
678 + hc_eps,
679 + hc_mult,
680 + )
681
682
683 direct_register_custom_op(
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+9
@@ -843,6 +843,15 @@ def w8a8_triton_block_scaled_mm(
843 assert len(block_size) == 2
844 block_n, block_k = block_size[0], block_size[1]
845
846 + # Triton cannot currently bind E8M0 scale tensors directly. On ROCm,
847 + # DeepSeek-V4 checkpoints store block scales in exponent-only E8M0 format,
848 + # so decode them to fp32 before launching the kernel.
849 + if current_platform.is_rocm():
850 + if As.dtype == torch.float8_e8m0fnu:
851 + As = _upcast_e8m0_to_fp32(As).contiguous()
852 + if Bs.dtype == torch.float8_e8m0fnu:
853 + Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
854 +
855 assert A.shape[-1] == B.shape[-1]
856 assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
857 assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
vllm/model_executor/layers/sparse_attn_indexer.py
+22 -8
@@ -499,13 +499,31 @@ class SparseAttnIndexer(CustomOp):
499 k: torch.Tensor,
500 weights: torch.Tensor,
501 ):
502 - assert not self.skip_k_cache_insert, (
503 - "AMD platform doesn't support skip cache insert yet"
504 - )
502 assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
503 assert isinstance(q_quant, torch.Tensor), (
504 "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
505 )
506 + if self.skip_k_cache_insert or not rocm_aiter_ops.is_enabled():
507 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
508 + rocm_aiter_sparse_attn_indexer_native,
509 + )
510 +
511 + return rocm_aiter_sparse_attn_indexer_native(
512 + hidden_states,
513 + _encode_layer_name(self.k_cache.prefix),
514 + self.k_cache.kv_cache,
515 + q_quant,
516 + k,
517 + weights,
518 + self.quant_block_size,
519 + self.scale_fmt,
520 + self.topk_tokens,
521 + self.head_dim,
522 + self.max_model_len,
523 + self.max_total_seq_len,
524 + self.topk_indices_buffer,
525 + skip_k_cache_insert=self.skip_k_cache_insert,
526 + )
527 if rocm_aiter_ops.is_enabled():
528 return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
529 hidden_states,
@@ -522,8 +540,4 @@ class SparseAttnIndexer(CustomOp):
540 self.max_total_seq_len,
541 self.topk_indices_buffer,
542 )
525 - else:
526 - raise RuntimeError(
527 - "Sparse attention indexer ROCm custom op requires ROCm "
528 - "Aiter ops to be enabled."
529 - )
543 + raise RuntimeError("Sparse attention indexer ROCm path could not be selected.")
vllm/model_executor/models/deepseek_v4.py
+6 -1
@@ -1245,7 +1245,12 @@ class DeepseekV4Model(nn.Module):
1245 # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute
1246 # (compressor kv_score, indexer.weights_proj, indexer.compressor
1247 # kv_score). fused_wqa_wkv stays on the default stream.
1248 - aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
1248 + # Disable them on ROCm because of hang issues.
1249 + aux_stream_list = (
1250 + None
1251 + if current_platform.is_rocm()
1252 + else [torch.cuda.Stream() for _ in range(3)]
1253 + )
1254
1255 self.device = current_platform.device_type
1256 # Reserved topk indices buffer for all Indexer layers to reuse.
vllm/model_executor/models/deepseek_v4_mtp.py
+6 -2
@@ -167,8 +167,12 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
167 )
168
169 # Three aux streams shared across all MTP layers, mirroring
170 - # DeepseekV4Model.
171 - aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
170 + # DeepseekV4Model. ROCm runs the same work serially for now.
171 + aux_stream_list = (
172 + None
173 + if current_platform.is_rocm()
174 + else [torch.cuda.Stream() for _ in range(3)]
175 + )
176
177 # to map the exact layer index from weights
178 self.layers = torch.nn.ModuleDict(
vllm/platforms/rocm.py
+1
@@ -409,6 +409,7 @@ class RocmPlatform(Platform):
409 "gptq",
410 "gptq_marlin", # will be overwritten with gptq
411 "fp8",
412 + "deepseek_v4_fp8",
413 "compressed-tensors",
414 "fbgemm_fp8",
415 "gguf",
vllm/v1/attention/backends/mla/sparse_swa.py
+2 -1
@@ -7,6 +7,7 @@ import torch
7
8 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
9 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10 +from vllm.platforms import current_platform
11 from vllm.triton_utils import tl, triton
12 from vllm.v1.attention.backend import (
13 AttentionBackend,
@@ -360,7 +361,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
361 _LAYER_TYPE_C4A: None,
362 _LAYER_TYPE_C128A: None,
363 }
363 - if num_decode_tokens == 0:
364 + if num_decode_tokens == 0 or current_platform.is_rocm():
365 return out
366 for layer_type in self._layer_types:
367 # get_mla_metadata() is the official FlashMLA entry point that
vllm/v1/attention/ops/deepseek_v4_ops/fused_inv_rope_fp8_quant.py
+3 -1
@@ -9,6 +9,7 @@ INT32-packed UE8M0 on SM100) so fp8_einsum skips transform_sf_into_required_layo
9
10 import torch
11
12 +from vllm.platforms import current_platform
13 from vllm.triton_utils import tl, triton
14 from vllm.utils.torch_utils import direct_register_custom_op
15
@@ -242,6 +243,7 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
243 (scale_inner * tma_aligned_T, 1, tma_aligned_T),
244 )
245 grid = (tma_aligned_T, n_groups * heads_per_group)
246 + pdl_kwargs = {} if current_platform.is_rocm() else {"launch_pdl": False}
247 _fused_inv_rope_fp8_quant_per_head[grid](
248 o,
249 positions,
@@ -265,7 +267,7 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
267 HALF_ROPE=half_rope,
268 TMA_ALIGNED_SCALES=tma_aligned_scales,
269 num_stages=1,
268 - launch_pdl=False,
270 + **pdl_kwargs,
271 num_warps=1,
272 )
273 return fp8_buf, scale_buf
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+528 -60
@@ -2,9 +2,11 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import functools
4 import importlib
5 +import math
6 from importlib.util import find_spec
7
8 import torch
9 +import torch.nn.functional as F
10
11 from vllm.forward_context import get_forward_context
12 from vllm.platforms import current_platform
@@ -13,6 +15,11 @@ from vllm.utils.torch_utils import LayerNameType
15 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
16 from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
17
18 +if current_platform.is_rocm():
19 + from vllm.platforms.rocm import _ON_GFX942
20 +else:
21 + _ON_GFX942 = False
22 +
23
24 @triton.jit
25 def _indexer_k_quant_and_cache_kernel(
@@ -230,6 +237,43 @@ def fp8_paged_mqa_logits_torch(
237
238 fp8_dtype = current_platform.fp8_dtype()
239 batch_size, next_n, _, dim = q.size()
240 + if next_n == 1:
241 + block_size = kv_cache.shape[1]
242 + logits = torch.full(
243 + [batch_size, max_model_len],
244 + float("-inf"),
245 + device=q.device,
246 + dtype=torch.float32,
247 + )
248 + if context_lens.dim() > 1:
249 + context_lens = context_lens.squeeze(-1)
250 + kv_cache_flat = kv_cache.view(-1, block_size * (dim + 4))
251 + for i in range(batch_size):
252 + q_i = q[i, 0].to(torch.float32)
253 + q_scale = weights[i]
254 + seq_len = int(context_lens[i].item())
255 + assert seq_len <= max_model_len
256 + num_pages = cdiv(seq_len, block_size)
257 + padded_seq_len = num_pages * block_size
258 + pages = block_tables[i, :num_pages]
259 + cache = kv_cache_flat[pages]
260 + scale_offset = block_size * dim
261 + cache_value = (
262 + cache[..., :scale_offset].view(dtype=fp8_dtype).to(torch.float32)
263 + )
264 + cache_scale = (
265 + cache[..., scale_offset:].view(dtype=torch.float32).contiguous()
266 + )
267 + cache_value = cache_value.view(padded_seq_len, dim)
268 + cache_scale = cache_scale.view(padded_seq_len)
269 + score = F.linear(cache_value, q_i)
270 + score = F.relu(score)
271 + score *= q_scale[None, :]
272 + score = score.sum(dim=1)
273 + score *= cache_scale
274 + logits[i, :seq_len] = score[:seq_len]
275 + return logits
276 +
277 kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:]
278 scale = scale.contiguous().view(torch.float)
279 q = q.float()
@@ -241,20 +285,30 @@ def fp8_paged_mqa_logits_torch(
285 device=q.device,
286 dtype=torch.float32,
287 )
244 - context_lens = context_lens.tolist()
288 for i in range(batch_size):
289 context_len = context_lens[i]
247 - q_offsets = torch.arange(context_len - next_n, context_len, device="cuda")
290 + if context_len.ndim == 0:
291 + context_len_i = int(context_len.item())
292 + q_offsets = torch.arange(
293 + context_len_i - next_n, context_len_i, device=q.device
294 + )
295 + context_limit = torch.full(
296 + (next_n,), context_len_i, dtype=torch.int32, device=q.device
297 + )
298 + else:
299 + context_limit = context_len.to(device=q.device, dtype=torch.int32)
300 + q_offsets = context_limit - 1
301 weight_slice = (
302 weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous()
303 )
251 - for block_rk in range(cdiv(context_len, block_size)):
304 + max_context_len = int(context_limit.max().item())
305 + for block_rk in range(cdiv(max_context_len, block_size)):
306 block_idx = block_tables[i][block_rk]
307 qx, kx = q[i], kv_cache[block_idx]
308 k_offsets = torch.arange(
255 - block_rk * block_size, (block_rk + 1) * block_size, device="cuda"
309 + block_rk * block_size, (block_rk + 1) * block_size, device=q.device
310 )
257 - mask = (k_offsets[None, :] < context_len) & (
311 + mask = (k_offsets[None, :] < context_limit[:, None]) & (
312 k_offsets[None, :] <= q_offsets[:, None]
313 )
314 s = torch.where(
@@ -331,30 +385,52 @@ def rocm_fp8_paged_mqa_logits(
385 aiter_paged_mqa_logits_module = paged_mqa_logits_module()
386
387 if aiter_paged_mqa_logits_module is not None:
334 - deepgemm_fp8_paged_mqa_logits = (
335 - aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits
388 + if _ON_GFX942:
389 + deepgemm_fp8_paged_mqa_logits = (
390 + aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits
391 + )
392 + batch_size, next_n, heads, _ = q_fp8.shape
393 + out_logits = torch.full(
394 + [batch_size * next_n, max_model_len],
395 + float("-inf"),
396 + device="cuda",
397 + dtype=torch.float32,
398 + )
399 + deepgemm_fp8_paged_mqa_logits(
400 + q_fp8,
401 + kv_cache_fp8,
402 + weights,
403 + out_logits,
404 + context_lens,
405 + block_tables,
406 + max_model_len,
407 + ChunkK=256,
408 + Preshuffle=block_size == 64,
409 + KVBlockSize=block_size,
410 + WavePerEU=2,
411 + )
412 + return out_logits
413 + deepgemm_fp8_paged_mqa_logits_stage1 = (
414 + aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1
415 )
416 batch_size, next_n, heads, _ = q_fp8.shape
338 - out_logits = torch.full(
339 - [batch_size * next_n, max_model_len],
417 + out_qk = torch.full(
418 + (heads, batch_size * next_n, max_model_len),
419 float("-inf"),
420 device="cuda",
421 dtype=torch.float32,
422 )
344 - deepgemm_fp8_paged_mqa_logits(
423 + deepgemm_fp8_paged_mqa_logits_stage1(
424 q_fp8,
425 kv_cache_fp8,
426 weights,
348 - out_logits,
427 + out_qk,
428 context_lens,
429 block_tables,
430 max_model_len,
352 - ChunkK=256,
353 - Preshuffle=block_size == 64,
354 - KVBlockSize=block_size,
355 - WavePerEU=2,
431 + ChunkQ=heads,
432 )
357 - return out_logits
433 + return out_qk.sum(dim=0)
434 else:
435 return fp8_paged_mqa_logits_torch(
436 q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
@@ -464,6 +540,27 @@ def rocm_fp8_mqa_logits(
540 return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
541
542
543 +def _topk_indices_torch(logits: torch.Tensor, topk_tokens: int) -> torch.Tensor:
544 + k = min(topk_tokens, logits.shape[-1])
545 + values, indices = torch.topk(logits, k=k, dim=-1)
546 + indices = indices.to(torch.int32)
547 + indices = torch.where(
548 + values == float("-inf"),
549 + torch.full_like(indices, -1, dtype=torch.int32),
550 + indices,
551 + )
552 + if k == topk_tokens:
553 + return indices
554 + padded = torch.full(
555 + (logits.shape[0], topk_tokens),
556 + -1,
557 + dtype=torch.int32,
558 + device=logits.device,
559 + )
560 + padded[:, :k] = indices
561 + return padded
562 +
563 +
564 def rocm_aiter_sparse_attn_indexer_fake(
565 hidden_states: torch.Tensor,
566 k_cache_prefix: LayerNameType,
@@ -482,8 +579,9 @@ def rocm_aiter_sparse_attn_indexer_fake(
579 # profile run
580 # NOTE(Chen): create the max possible flattened_kv. So that
581 # profile_run can get correct memory usage.
582 + device = hidden_states.device if k is None else k.device
583 _flattened_kv = torch.empty(
486 - [total_seq_lens, head_dim + 4], device=k.device, dtype=torch.uint8
584 + [total_seq_lens, head_dim + 4], device=device, dtype=torch.uint8
585 )
586 fp8_dtype = current_platform.fp8_dtype()
587 _k_fp8 = _flattened_kv[..., :head_dim].view(fp8_dtype).contiguous()
@@ -491,7 +589,7 @@ def rocm_aiter_sparse_attn_indexer_fake(
589 return topk_indices_buffer
590
591
494 -def rocm_aiter_sparse_attn_indexer(
592 +def rocm_aiter_sparse_attn_indexer_native(
593 hidden_states: torch.Tensor,
594 k_cache_prefix: LayerNameType,
595 kv_cache: torch.Tensor,
@@ -505,10 +603,12 @@ def rocm_aiter_sparse_attn_indexer(
603 max_model_len: int,
604 total_seq_lens: int,
605 topk_indices_buffer: torch.Tensor | None,
606 + skip_k_cache_insert: bool = False,
607 ) -> torch.Tensor:
608 # careful! this will be None in dummy run
609 attn_metadata = get_forward_context().attn_metadata
610 fp8_dtype = current_platform.fp8_dtype()
611 + from vllm import _custom_ops as ops
612 from vllm.utils.torch_utils import _resolve_layer_name
613
614 k_cache_prefix = _resolve_layer_name(k_cache_prefix)
@@ -537,19 +637,33 @@ def rocm_aiter_sparse_attn_indexer(
637 has_decode = layer_attn_metadata.num_decodes > 0
638 has_prefill = layer_attn_metadata.num_prefills > 0
639 num_decode_tokens = layer_attn_metadata.num_decode_tokens
640 + device = hidden_states.device if k is None else k.device
641
642 # during speculative decoding, k may be padded to the CUDA graph batch
643 # size while slot_mapping only covers actual tokens.
644 num_tokens = slot_mapping.shape[0]
544 - k = k[:num_tokens]
645 + if k is not None:
646 + k = k[:num_tokens]
647 + elif not skip_k_cache_insert:
648 + raise ValueError("k must be provided when skip_k_cache_insert is False")
649
546 - indexer_k_quant_and_cache_triton(
547 - k,
548 - kv_cache,
549 - slot_mapping,
550 - quant_block_size,
551 - scale_fmt,
552 - )
650 + if not skip_k_cache_insert:
651 + if _ON_GFX942:
652 + ops.indexer_k_quant_and_cache(
653 + k,
654 + kv_cache,
655 + slot_mapping,
656 + quant_block_size,
657 + scale_fmt,
658 + )
659 + else:
660 + indexer_k_quant_and_cache_triton(
661 + k,
662 + kv_cache,
663 + slot_mapping,
664 + quant_block_size,
665 + scale_fmt,
666 + )
667
668 topk_indices_buffer[: hidden_states.shape[0]] = -1
669 if has_prefill:
@@ -558,22 +672,31 @@ def rocm_aiter_sparse_attn_indexer(
672 for chunk in prefill_metadata.chunks:
673 k_fp8 = torch.empty(
674 [chunk.total_seq_lens, head_dim],
561 - device=k.device,
675 + device=device,
676 dtype=fp8_dtype,
677 )
678 k_scale = torch.empty(
679 [chunk.total_seq_lens, 4],
566 - device=k.device,
680 + device=device,
681 dtype=torch.uint8,
682 )
569 - cp_gather_indexer_k_quant_cache_triton(
570 - kv_cache,
571 - k_fp8,
572 - k_scale,
573 - chunk.block_table,
574 - chunk.cu_seq_lens,
575 - chunk.token_to_seq,
576 - )
683 + if _ON_GFX942:
684 + ops.cp_gather_indexer_k_quant_cache(
685 + kv_cache,
686 + k_fp8,
687 + k_scale,
688 + chunk.block_table,
689 + chunk.cu_seq_lens,
690 + )
691 + else:
692 + cp_gather_indexer_k_quant_cache_triton(
693 + kv_cache,
694 + k_fp8,
695 + k_scale,
696 + chunk.block_table,
697 + chunk.cu_seq_lens,
698 + token_to_seq=chunk.token_to_seq,
699 + )
700
701 logits = rocm_fp8_mqa_logits(
702 q_fp8[chunk.token_start : chunk.token_end],
@@ -582,21 +705,10 @@ def rocm_aiter_sparse_attn_indexer(
705 chunk.cu_seqlen_ks,
706 chunk.cu_seqlen_ke,
707 )
585 - num_rows = logits.shape[0]
586 - assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
708 topk_indices = topk_indices_buffer[
709 chunk.token_start : chunk.token_end, :topk_tokens
710 ]
590 - torch.ops._C.top_k_per_row_prefill(
591 - logits,
592 - chunk.cu_seqlen_ks,
593 - chunk.cu_seqlen_ke,
594 - topk_indices,
595 - num_rows,
596 - logits.stride(0),
597 - logits.stride(1),
598 - topk_tokens,
599 - )
711 + topk_indices.copy_(_topk_indices_torch(logits, topk_tokens))
712
713 if has_decode:
714 decode_metadata = layer_attn_metadata.decode
@@ -633,19 +745,8 @@ def rocm_aiter_sparse_attn_indexer(
745 max_model_len=max_model_len,
746 )
747
636 - num_rows = logits.shape[0]
637 - assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
748 topk_indices = topk_indices_buffer[:num_decode_tokens, :topk_tokens]
639 - torch.ops._C.top_k_per_row_decode(
640 - logits,
641 - next_n,
642 - decode_metadata.seq_lens,
643 - topk_indices,
644 - num_rows,
645 - logits.stride(0),
646 - logits.stride(1),
647 - topk_tokens,
648 - )
749 + topk_indices.copy_(_topk_indices_torch(logits, topk_tokens)[:num_decode_tokens])
750
751 if decode_metadata.requires_padding:
752 # if padded, we need to unpack
@@ -659,3 +760,370 @@ def rocm_aiter_sparse_attn_indexer(
760 )
761
762 return topk_indices_buffer
763 +
764 +
765 +def rocm_aiter_sparse_attn_indexer(
766 + hidden_states: torch.Tensor,
767 + k_cache_prefix: LayerNameType,
768 + kv_cache: torch.Tensor,
769 + q_fp8: torch.Tensor,
770 + k: torch.Tensor,
771 + weights: torch.Tensor,
772 + quant_block_size: int,
773 + scale_fmt: str | None,
774 + topk_tokens: int,
775 + head_dim: int,
776 + max_model_len: int,
777 + total_seq_lens: int,
778 + topk_indices_buffer: torch.Tensor | None,
779 +) -> torch.Tensor:
780 + return rocm_aiter_sparse_attn_indexer_native(
781 + hidden_states,
782 + k_cache_prefix,
783 + kv_cache,
784 + q_fp8,
785 + k,
786 + weights,
787 + quant_block_size,
788 + scale_fmt,
789 + topk_tokens,
790 + head_dim,
791 + max_model_len,
792 + total_seq_lens,
793 + topk_indices_buffer,
794 + skip_k_cache_insert=False,
795 + )
796 +
797 +
798 +def _decode_e8m0_scales(scale: torch.Tensor) -> torch.Tensor:
799 + if scale.dtype == torch.float8_e8m0fnu:
800 + from vllm.model_executor.layers.quantization.utils.fp8_utils import (
801 + _upcast_e8m0_to_fp32,
802 + )
803 +
804 + return _upcast_e8m0_to_fp32(scale).contiguous()
805 + return scale.to(torch.float32)
806 +
807 +
808 +def _expand_2d_block_scales(
809 + scale: torch.Tensor,
810 + rows: int,
811 + cols: int,
812 +) -> torch.Tensor:
813 + scale = _decode_e8m0_scales(scale)
814 + row_blocks, col_blocks = scale.shape[-2:]
815 + row_block = math.ceil(rows / row_blocks)
816 + col_block = math.ceil(cols / col_blocks)
817 + scale = torch.repeat_interleave(scale, row_block, dim=-2)[..., :rows, :]
818 + scale = torch.repeat_interleave(scale, col_block, dim=-1)[..., :, :cols]
819 + return scale
820 +
821 +
822 +def _apply_gptj_inv_rope_ref(
823 + x: torch.Tensor,
824 + positions: torch.Tensor,
825 + cos_sin_cache: torch.Tensor,
826 + rope_dim: int,
827 +) -> torch.Tensor:
828 + if rope_dim == 0 or x.numel() == 0:
829 + return x
830 + half_rot = rope_dim // 2
831 + nope_dim = x.shape[-1] - rope_dim
832 + dtype = x.dtype
833 + x = x.to(torch.float32)
834 + cache = cos_sin_cache.index_select(0, positions.to(torch.long))
835 + cos = cache[:, :half_rot].to(torch.float32)
836 + sin = cache[:, half_rot : 2 * half_rot].to(torch.float32)
837 + view_shape = (positions.shape[0],) + (1,) * (x.dim() - 2) + (half_rot,)
838 + cos = cos.view(view_shape)
839 + sin = sin.view(view_shape)
840 + rope = x[..., nope_dim:]
841 + y_even = rope[..., 0::2]
842 + y_odd = rope[..., 1::2]
843 + rope_out = torch.stack(
844 + (y_even * cos + y_odd * sin, y_odd * cos - y_even * sin),
845 + dim=-1,
846 + ).flatten(-2)
847 + x = x.clone()
848 + x[..., nope_dim:] = rope_out
849 + return x.to(dtype)
850 +
851 +
852 +def _apply_inv_rope_ref(
853 + rotary_emb: torch.nn.Module,
854 + x: torch.Tensor,
855 + positions: torch.Tensor,
856 + rope_dim: int,
857 +) -> torch.Tensor:
858 + if hasattr(rotary_emb, "forward_native"):
859 + try:
860 + query, _ = rotary_emb.forward_native(
861 + positions,
862 + x.clone(),
863 + None,
864 + inverse=True,
865 + )
866 + return query
867 + except TypeError:
868 + pass
869 + return _apply_gptj_inv_rope_ref(x, positions, rotary_emb.cos_sin_cache, rope_dim)
870 +
871 +
872 +def rocm_inv_rope_einsum(
873 + rotary_emb: torch.nn.Module,
874 + o: torch.Tensor,
875 + positions: torch.Tensor,
876 + rope_head_dim: int,
877 + n_local_groups: int,
878 + o_lora_rank: int,
879 + wo_a: torch.nn.Module,
880 +) -> torch.Tensor:
881 + """Reference inverse-RoPE + WO_A einsum path used on ROCm."""
882 + o_ref = _apply_inv_rope_ref(rotary_emb, o, positions, rope_head_dim).to(
883 + torch.bfloat16
884 + )
885 + o_ref = o_ref.view(o.shape[0], n_local_groups, -1)
886 +
887 + hidden_dim = o_ref.shape[-1]
888 + if hasattr(wo_a, "weight_scale_inv"):
889 + wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
890 + torch.float32
891 + )
892 + wo_a_scale = _expand_2d_block_scales(
893 + wo_a.weight_scale_inv.view(
894 + n_local_groups, -1, wo_a.weight_scale_inv.shape[-1]
895 + ),
896 + o_lora_rank,
897 + hidden_dim,
898 + )
899 + wo_a_weight = (wo_a_weight * wo_a_scale).to(torch.bfloat16)
900 + else:
901 + wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
902 + torch.bfloat16
903 + )
904 +
905 + return torch.einsum("tgd,grd->tgr", o_ref, wo_a_weight)
906 +
907 +
908 +def rocm_ref_sparse_attn_prefill(
909 + q: torch.Tensor,
910 + kv: torch.Tensor,
911 + indices: torch.Tensor,
912 + topk_length: torch.Tensor | None,
913 + scale: float,
914 + head_dim: int,
915 + attn_sink: torch.Tensor | None,
916 +) -> torch.Tensor:
917 + indices = indices.clone().squeeze(1)
918 + s_q, h_q, d_qk = q.shape
919 + topk = indices.shape[-1]
920 + s_kv = kv.shape[0]
921 + if topk_length is not None:
922 + mask = torch.arange(topk, device=indices.device).unsqueeze(
923 + 0
924 + ) >= topk_length.unsqueeze(1)
925 + indices[mask] = -1
926 + invalid_mask = (indices < 0) | (indices >= s_kv)
927 + indices[invalid_mask] = 0
928 +
929 + qf = q.float()
930 + gathered_kv = kv.index_select(0, indices.flatten()).reshape(s_q, topk, d_qk).float()
931 + scores = qf @ gathered_kv.transpose(1, 2)
932 + scores *= scale
933 + scores[invalid_mask.unsqueeze(1).expand_as(scores)] = float("-inf")
934 +
935 + orig_lse = torch.logsumexp(scores, dim=-1)
936 + lse_for_o = orig_lse
937 + if attn_sink is not None:
938 + lse_for_o = torch.logsumexp(
939 + torch.stack(
940 + [orig_lse, attn_sink[:h_q].view(1, h_q).expand_as(orig_lse)],
941 + dim=0,
942 + ),
943 + dim=0,
944 + )
945 + lse_for_o = lse_for_o.clone()
946 + lse_for_o[lse_for_o == float("-inf")] = float("+inf")
947 + probs = torch.exp(scores - lse_for_o.unsqueeze(-1))
948 + out = probs @ gathered_kv[..., :head_dim]
949 + lonely_q_mask = orig_lse == float("-inf")
950 + out[lonely_q_mask.unsqueeze(-1).expand_as(out)] = 0.0
951 + return out.to(torch.bfloat16)
952 +
953 +
954 +def rocm_sparse_attn_prefill(
955 + q: torch.Tensor,
956 + kv: torch.Tensor,
957 + indices: torch.Tensor,
958 + topk_length: torch.Tensor | None,
959 + scale: float,
960 + head_dim: int,
961 + attn_sink: torch.Tensor | None,
962 + output: torch.Tensor,
963 +) -> None:
964 + output_chunk = rocm_ref_sparse_attn_prefill(
965 + q=q,
966 + kv=kv,
967 + indices=indices,
968 + topk_length=topk_length,
969 + scale=scale,
970 + head_dim=head_dim,
971 + attn_sink=attn_sink,
972 + )
973 + output.copy_(output_chunk.to(output.dtype))
974 +
975 +
976 +def rocm_dequantize_blocked_k_cache(
977 + quant_k_cache: torch.Tensor,
978 + head_dim: int,
979 + nope_head_dim: int,
980 + rope_head_dim: int,
981 +) -> torch.Tensor:
982 + fp8_dtype = current_platform.fp8_dtype()
983 + tile_size = 64
984 + num_tiles = nope_head_dim // tile_size
985 +
986 + num_blocks, block_size, _ = quant_k_cache.shape
987 + quant_k_cache = quant_k_cache.view(num_blocks, -1)
988 + input_nope_rope = quant_k_cache[
989 + :, : block_size * (nope_head_dim + 2 * rope_head_dim)
990 + ].view(num_blocks, block_size, nope_head_dim + 2 * rope_head_dim)
991 + input_nope = input_nope_rope[:, :, :nope_head_dim].view(fp8_dtype)
992 + input_rope = input_nope_rope[:, :, nope_head_dim:].view(torch.bfloat16)
993 + input_scale = (
994 + quant_k_cache[:, block_size * (nope_head_dim + 2 * rope_head_dim) :]
995 + .view(num_blocks, block_size, 8)[:, :, :num_tiles]
996 + .view(torch.float8_e8m0fnu)
997 + )
998 +
999 + result = torch.empty(
1000 + (num_blocks, block_size, 1, head_dim),
1001 + dtype=torch.bfloat16,
1002 + device=quant_k_cache.device,
1003 + )
1004 + result[..., nope_head_dim:] = input_rope.unsqueeze(2)
1005 + for tile_idx in range(num_tiles):
1006 + cur_nope = input_nope[
1007 + ..., tile_idx * tile_size : (tile_idx + 1) * tile_size
1008 + ].to(torch.bfloat16)
1009 + cur_scales = input_scale[:, :, tile_idx].to(torch.bfloat16).unsqueeze(-1)
1010 + result[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
1011 + cur_nope * cur_scales
1012 + ).unsqueeze(2)
1013 + return result
1014 +
1015 +
1016 +def rocm_ref_sparse_attn_decode(
1017 + q: torch.Tensor,
1018 + blocked_k: torch.Tensor,
1019 + indices_in_kvcache: torch.Tensor,
1020 + topk_length: torch.Tensor | None,
1021 + scale: float,
1022 + head_dim: int,
1023 + attn_sink: torch.Tensor | None,
1024 + extra_blocked_k: torch.Tensor | None = None,
1025 + extra_indices_in_kvcache: torch.Tensor | None = None,
1026 + extra_topk_length: torch.Tensor | None = None,
1027 +) -> torch.Tensor:
1028 + b, s_q, h_q, d_qk = q.shape
1029 +
1030 + def process_scope(
1031 + cur_blocked_k: torch.Tensor,
1032 + cur_indices: torch.Tensor,
1033 + cur_topk_length: torch.Tensor | None,
1034 + ) -> tuple[torch.Tensor, torch.Tensor]:
1035 + cur_indices = cur_indices.reshape(b, s_q, -1)
1036 + topk = cur_indices.size(-1)
1037 + fixed_indices = torch.clamp_min(cur_indices, 0)
1038 + gathered_kv = (
1039 + cur_blocked_k.view(-1, d_qk)
1040 + .index_select(0, fixed_indices.view(-1))
1041 + .view(b, s_q, topk, d_qk)
1042 + )
1043 + invalid_mask = cur_indices == -1
1044 + if cur_topk_length is not None:
1045 + cur_topk_length = cur_topk_length.reshape(b)
1046 + invalid_mask |= torch.arange(0, topk, device=invalid_mask.device).view(
1047 + 1, 1, topk
1048 + ) >= cur_topk_length.view(b, 1, 1)
1049 + return gathered_kv, invalid_mask
1050 +
1051 + gathered_kv, invalid_mask = process_scope(
1052 + blocked_k, indices_in_kvcache, topk_length
1053 + )
1054 + if extra_blocked_k is not None:
1055 + assert extra_indices_in_kvcache is not None
1056 + gathered_kv1, invalid_mask1 = process_scope(
1057 + extra_blocked_k, extra_indices_in_kvcache, extra_topk_length
1058 + )
1059 + gathered_kv = torch.cat([gathered_kv, gathered_kv1], dim=2)
1060 + invalid_mask = torch.cat([invalid_mask, invalid_mask1], dim=2)
1061 +
1062 + gathered_kv = gathered_kv.view(b * s_q, -1, d_qk).float()
1063 + gathered_kv[gathered_kv != gathered_kv] = 0.0
1064 + qf = q.float().view(b * s_q, h_q, d_qk)
1065 + attn_weight = qf @ gathered_kv.transpose(-1, -2)
1066 + attn_weight *= scale
1067 + attn_weight[
1068 + invalid_mask.view(b * s_q, 1, -1).expand(b * s_q, h_q, invalid_mask.size(-1))
1069 + ] = float("-inf")
1070 + lse = attn_weight.logsumexp(dim=-1)
1071 + attn_weight = torch.exp(attn_weight - lse.unsqueeze(-1))
1072 + output = attn_weight @ gathered_kv[..., :head_dim]
1073 + output = output.view(b, s_q, h_q, head_dim)
1074 + lse = lse.view(b, s_q, h_q)
1075 +
1076 + if attn_sink is not None:
1077 + output *= (1.0 / (1.0 + torch.exp(attn_sink.view(1, 1, h_q) - lse))).unsqueeze(
1078 + -1
1079 + )
1080 +
1081 + lonely_q_mask = lse == float("-inf")
1082 + output[lonely_q_mask.unsqueeze(-1).expand_as(output)] = 0.0
1083 + return output.squeeze(1).to(torch.bfloat16)
1084 +
1085 +
1086 +def rocm_forward_decode_fallback(
1087 + q: torch.Tensor,
1088 + kv_cache: torch.Tensor | None,
1089 + swa_k_cache: torch.Tensor,
1090 + swa_only: bool,
1091 + topk_indices: torch.Tensor | None,
1092 + topk_lens: torch.Tensor | None,
1093 + swa_indices: torch.Tensor,
1094 + swa_lens: torch.Tensor,
1095 + attn_sink: torch.Tensor | None,
1096 + scale: float,
1097 + head_dim: int,
1098 + nope_head_dim: int,
1099 + rope_head_dim: int,
1100 + output: torch.Tensor,
1101 +) -> None:
1102 + blocked_swa = rocm_dequantize_blocked_k_cache(
1103 + swa_k_cache,
1104 + head_dim=head_dim,
1105 + nope_head_dim=nope_head_dim,
1106 + rope_head_dim=rope_head_dim,
1107 + )
1108 + blocked_extra = None
1109 + if not swa_only:
1110 + assert kv_cache is not None
1111 + blocked_extra = rocm_dequantize_blocked_k_cache(
1112 + kv_cache,
1113 + head_dim=head_dim,
1114 + nope_head_dim=nope_head_dim,
1115 + rope_head_dim=rope_head_dim,
1116 + )
1117 + attn_out = rocm_ref_sparse_attn_decode(
1118 + q=q.unsqueeze(1),
1119 + blocked_k=blocked_swa,
1120 + indices_in_kvcache=swa_indices.unsqueeze(1),
1121 + topk_length=swa_lens,
1122 + scale=scale,
1123 + head_dim=head_dim,
1124 + attn_sink=attn_sink[: q.shape[1]] if attn_sink is not None else None,
1125 + extra_blocked_k=blocked_extra,
1126 + extra_indices_in_kvcache=topk_indices,
1127 + extra_topk_length=topk_lens,
1128 + )
1129 + output.copy_(attn_out.to(output.dtype))