@tduncan / Nvidia-Optimization-Stack-VLLM / commits / afdcbd5d39

[ROCm][DSv4] Functional fixes for DeepSeek V4 on MI300X/MI325X (#45681)

Signed-off-by: ganyi <ygan@amd.com> Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com> Co-authored-by: ganyi <ygan@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Jin Tao <jintao12@amd.com>

Tuukka Sarvi committed Jun 18, 2026 at 15:21 UTC afdcbd5d39eaf2b37b616c8ee8aabc51e15e70ef

8 files changed +545 -52

csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

+11 -4

@@ -18,7 +18,7 @@

18 * ROPE_DIM = 64 (RoPE applied to dims [NOPE_DIM, HEAD_DIM))

19 * NOPE_DIM = 448

20 * QUANT_BLOCK = 64 (UE8M0 FP8 quant block)

21 - * FP8_MAX = 448.0f

21 + * FP8_MAX = 224.0f on ROCm FNUZ / 448.0f on OCP

22 * is_neox=false (GPT-J interleaved pairs)

23 * cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin

24 * second along last dim; each half is rope_dim/2 = 32 values)

@@ -61,10 +61,11 @@

61 #ifdef USE_ROCM

62 // ROCm-compatible FP8 conversion helpers

63 __device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {

64 - #if defined(HIP_FP8_TYPE_OCP)

65 - __hip_fp8_e4m3 fp8_val(val);

66 - #else

64 + // gfx942 uses FNUZ FP8; other ROCm targets use OCP E4M3.

65 + #if defined(__gfx942__)

66 __hip_fp8_e4m3_fnuz fp8_val(val);

67 + #else

68 + __hip_fp8_e4m3 fp8_val(val);

69 #endif

70 return reinterpret_cast<uint8_t&>(fp8_val);

71 }

@@ -90,7 +91,13 @@ constexpr int kQuantBlock = 64;

91 constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7

92 constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad)

93 constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576

94 +// FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw

95 +// 240.0) to match the rest of vLLM's FNUZ pipeline.

96 +#if defined(USE_ROCM) && defined(__gfx942__)

97 +constexpr float kFp8Max = 224.0f;

98 +#else

99 constexpr float kFp8Max = 448.0f;

100 +#endif

101

102 #ifndef USE_ROCM

103 // When num_tokens is less than this threshold,

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

+147 -24

@@ -19,17 +19,28 @@ The kernel is imported via

19 import pytest

20 import torch

22 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

23 + get_fp8_min_max,

24 +)

25 from vllm.models.deepseek_v4.common.ops import (

26 dequantize_and_gather_k_cache,

27 quantize_and_insert_k_cache,

28 )

29 +from vllm.platforms import current_platform

31 # ── Constants matching the kernel ────────────────────────────────────────────

32 HEAD_DIM = 512

33 ROPE_DIM = 64

34 NOPE_DIM = HEAD_DIM - ROPE_DIM # 448

35 QUANT_BLOCK = 64

32 -FP8_MAX = 448.0

36 +# Match the C++ SWA-K encoder: FNUZ on gfx942, OCP elsewhere.

37 +USE_FNUZ = current_platform.is_fp8_fnuz()

38 +_, FP8_MAX = get_fp8_min_max()

39 +# The kernel emits FNUZ-encoded fp8 bytes on gfx942 (rocm_cvt_float_to_fp8_e4m3)

40 +# but stores them into float8_e4m3fn-typed tensors, matching vLLM's ROCm cache

41 +# convention. References must encode under the same scheme and the kernel's

42 +# e4m3fn-typed outputs must be reinterpreted under it before decoding.

43 +FP8_STORE_DTYPE = torch.float8_e4m3fnuz if USE_FNUZ else torch.float8_e4m3fn

44 HEAD_BYTES = NOPE_DIM + ROPE_DIM * 2 + 8 # 448 + 128 + 8 = 584

@@ -81,10 +92,11 @@ def apply_rope_gptj_last_k(

92 cos = cos.unsqueeze(1)

93 sin = sin.unsqueeze(1)

84 - # Use addcmul (compiles to FMA on CUDA) for the 2x2 rotation. nvcc lowers

85 - # the kernel's `e*c - o*s` to fma(e, c, -o*s); matching that here keeps

86 - # near-cancellation pairs on the same bf16 grid as the kernel output and

87 - # avoids spurious 1-ULP boundary flips at high num_tokens.

95 + # Use addcmul (an FMA) for the 2x2 rotation to mirror the kernel's

96 + # `e*c - o*s` fused form. This keeps the reference close to the kernel, but

97 + # the fp32 reference and the fp32 GPU kernel can still round to bf16 on

98 + # opposite sides of a round-to-nearest tie for a tiny number of elements at

99 + # high positions, so callers compare the RoPE region within 1 bf16 ULP.

100 new_even = torch.addcmul(-odd * sin, even, cos)

101 new_odd = torch.addcmul(odd * cos, even, sin)

102 rope_rotated = torch.stack((new_even, new_odd), dim=-1).reshape(shape)

@@ -148,6 +160,86 @@ def _call_fused(

160 )

161

162

163 +def _bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:

164 + """Representable-step distance between two bf16 tensors.

165 +

166 + Reinterprets the bf16 bit patterns under the IEEE-754 total ordering so

167 + that adjacent representable values differ by exactly 1.

168 + """

169 +

170 + def key(t: torch.Tensor) -> torch.Tensor:

171 + u = t.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF

172 + return torch.where(u >= 0x8000, 0xFFFF - u, u + 0x8000)

173 +

174 + return (key(a) - key(b)).abs()

175 +

176 +

177 +def _fp8_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:

178 + """Representable-step distance between two 8-bit fp8 tensors.

179 +

180 + Reinterprets the fp8 bytes under a sign-magnitude total ordering so that

181 + adjacent representable values differ by exactly 1. Inputs must already share

182 + the same fp8 encoding (e.g. both FP8_STORE_DTYPE).

183 + """

184 +

185 + def key(t: torch.Tensor) -> torch.Tensor:

186 + u = t.contiguous().view(torch.uint8).to(torch.int64)

187 + return torch.where(u >= 0x80, 0xFF - u, u + 0x80)

188 +

189 + return (key(a) - key(b)).abs()

190 +

191 +

192 +def _as_stored_fp8(t: torch.Tensor) -> torch.Tensor:

193 + """Reinterpret a float8_e4m3fn-typed kernel output under the real (FNUZ on

194 + gfx942) encoding the kernel actually wrote, without touching the bytes."""

195 + return t.contiguous().view(torch.uint8).view(FP8_STORE_DTYPE)

196 +

197 +

198 +def _dequant_cache(k_cache_2d, num_tokens, num_blocks, block_size):

199 + """Round-trip a [num_blocks, block_size*HEAD_BYTES] K-cache back to bf16."""

200 + device = k_cache_2d.device

201 + out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device)

202 + seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device)

203 + block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(

204 + 0

205 + )

206 + k_cache_3d = k_cache_2d.view(num_blocks, block_size, HEAD_BYTES)

207 + dequantize_and_gather_k_cache(

208 + out,

209 + k_cache_3d,

210 + seq_lens,

211 + None,

212 + block_table,

213 + block_size,

214 + offset=0,

215 + use_fnuz=USE_FNUZ,

216 + )

217 + return out[0, :num_tokens]

218 +

219 +

220 +def _assert_kv_cache_parity(

221 + k_cache_fused, k_cache_ref, num_tokens, num_blocks, block_size

222 +):

223 + """Assert the fused and reference K-caches agree after decoding.

224 +

225 + The NoPE region is deterministic UE8M0 FP8, so its round-trip must be

226 + bit-identical. The RoPE region is stored as bf16 after an fp32 rotation:

227 + the GPU kernel and the PyTorch reference can fall on opposite sides of a

228 + round-to-nearest tie and differ by at most one bf16 ULP. (Spot checks show

229 + the kernel value is the correctly-rounded one; the fp32 torch reference is

230 + the one that lands on the wrong side near a midpoint.) Allow <=1 ULP there.

231 + """

232 + rec_fused = _dequant_cache(k_cache_fused, num_tokens, num_blocks, block_size)

233 + rec_ref = _dequant_cache(k_cache_ref, num_tokens, num_blocks, block_size)

234 + torch.testing.assert_close(

235 + rec_fused[:, :NOPE_DIM], rec_ref[:, :NOPE_DIM], rtol=0, atol=0

236 + )

237 + max_ulp = int(

238 + _bf16_ulp_distance(rec_fused[:, NOPE_DIM:], rec_ref[:, NOPE_DIM:]).max().item()

239 + )

240 + assert max_ulp <= 1, f"RoPE bf16 region differs by {max_ulp} ULP (>1)"

241 +

242 +

243 # ── Test 1: Q path numerical parity ──────────────────────────────────────────

244

245

@@ -241,7 +333,7 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int):

333 num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device

334 )

335 quantize_and_insert_k_cache(

244 - kv_ref, k_cache_ref, slot_mapping, block_size=block_size

336 + kv_ref, k_cache_ref, slot_mapping, block_size=block_size, use_fnuz=USE_FNUZ

337 )

338

339 # ── Fused path (dummy q, padded to FlashMLA's min head count 64) ───────

@@ -273,7 +365,14 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int):

365 # gather_lens arg is None (use seq_lens)

366 k_cache_3d = k_cache_2d.view(num_blocks, block_size, HEAD_BYTES)

367 dequantize_and_gather_k_cache(

276 - out, k_cache_3d, seq_lens, None, block_table, block_size, offset=0

368 + out,

369 + k_cache_3d,

370 + seq_lens,

371 + None,

372 + block_table,

373 + block_size,

374 + offset=0,

375 + use_fnuz=USE_FNUZ,

376 )

377 return out[0, :num_tokens]

378

@@ -297,12 +396,10 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int):

396 f"fused NoPE token {t} diff {diff_fused} > {max_allowed}"

397 )

398

300 - # RoPE region: bf16 stored exactly → zero diff.

301 - rope_diff = (recovered_fused[:, NOPE_DIM:] - kv_ref[:, NOPE_DIM:]).abs().max()

302 - assert rope_diff.item() == 0.0, f"RoPE portion not exact: {rope_diff.item()}"

303 -

304 - # Exact byte equality of the two cache buffers — strong parity.

305 - torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)

399 + # Strong parity: NoPE FP8 round-trip bit-identical, RoPE bf16 within 1 ULP.

400 + _assert_kv_cache_parity(

401 + k_cache_fused, k_cache_ref, num_tokens, num_blocks, block_size

402 + )

403

404

405 # ── Test 2b: DP padding (slot_mapping shorter than q/kv) ─────────────────────

@@ -336,7 +433,7 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):

433 num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device

434 )

435 quantize_and_insert_k_cache(

339 - kv_ref, k_cache_ref, slot_mapping, block_size=block_size

436 + kv_ref, k_cache_ref, slot_mapping, block_size=block_size, use_fnuz=USE_FNUZ

437 )

438

439 # Fused: pass full-sized q/kv/positions, shorter slot_mapping.

@@ -354,7 +451,9 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):

451 block_size,

452 )

453

357 - torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)

454 + _assert_kv_cache_parity(

455 + k_cache_fused, k_cache_ref, num_tokens, num_blocks, block_size

456 + )

457

458

459 # ── Test 3: combined single-call Q + KV parity ───────────────────────────────

@@ -403,7 +502,7 @@ def test_combined_q_and_kv(

502 num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device

503 )

504 quantize_and_insert_k_cache(

406 - kv_ref, k_cache_ref, slot_mapping, block_size=block_size

505 + kv_ref, k_cache_ref, slot_mapping, block_size=block_size, use_fnuz=USE_FNUZ

506 )

507

508 # Fused single call.

@@ -426,7 +525,9 @@ def test_combined_q_and_kv(

525 assert pad_region.abs().max().item() == 0.0, (

526 "padded head slots must be exact zero"

527 )

429 - torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)

528 + _assert_kv_cache_parity(

529 + k_cache_fused, k_cache_ref, num_tokens, num_blocks, block_size

530 + )

531

532

533 # ── Full-cache (FlashInfer) path parity ──────────────────────────────────────

@@ -499,7 +600,7 @@ def _fp8_full_cache_reference(

600 q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache)

601 q_fp8.copy_(

602 torch.clamp(q_ref.float() * q_fp8_scale_inv, -FP8_MAX, FP8_MAX).to(

502 - torch.float8_e4m3fn

603 + FP8_STORE_DTYPE

604 )

605 )

606

@@ -510,7 +611,7 @@ def _fp8_full_cache_reference(

611 pos_in_block = slots % block_size

612 k_cache[block_idx, pos_in_block] = torch.clamp(

613 kv_ref[valid].float() / fp8_scale, -FP8_MAX, FP8_MAX

513 - ).to(torch.float8_e4m3fn)

614 + ).to(FP8_STORE_DTYPE)

615

616

617 def _bf16_full_cache_reference(

@@ -565,12 +666,17 @@ def test_full_cache_per_tensor_fp8_matches_reference(

666 fp8_scale = torch.tensor([1.0], dtype=torch.float32, device=device)

667 q_fp8_scale_inv = torch.tensor([1.0], dtype=torch.float32, device=device)

668

568 - q_fp8_ref = torch.empty_like(q, dtype=torch.float8_e4m3fn)

669 + # References are encoded under the scheme the kernel actually writes

670 + # (FNUZ on gfx942); the kernel's own outputs must stay float8_e4m3fn-typed

671 + # because the op asserts that dtype.

672 + q_fp8_ref = torch.empty_like(q, dtype=FP8_STORE_DTYPE)

673 q_fp8_fused = torch.empty_like(q, dtype=torch.float8_e4m3fn)

674 k_cache_ref = torch.zeros(

675 + num_blocks, block_size, HEAD_DIM, dtype=FP8_STORE_DTYPE, device=device

676 + )

677 + k_cache_fused = torch.zeros(

678 num_blocks, block_size, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device

679 )

573 - k_cache_fused = torch.zeros_like(k_cache_ref)

680

681 _fp8_full_cache_reference(

682 q,

@@ -599,12 +705,29 @@ def test_full_cache_per_tensor_fp8_matches_reference(

705 block_size,

706 )

707

708 + # Q is RMSNorm(no-weight)+RoPE in fp32 before fp8 quant; the RMSNorm

709 + # reduction and RoPE rotation can land the kernel and the torch reference on

710 + # opposite sides of an fp8 round-to-nearest tie, so allow <=1 fp8 ULP.

711 + q_fused = _as_stored_fp8(q_fp8_fused)

712 + q_max_ulp = int(_fp8_ulp_distance(q_fused, q_fp8_ref).max().item())

713 + assert q_max_ulp <= 1, f"Q fp8 differs by {q_max_ulp} ULP (>1)"

714 +

715 + # K-cache NoPE region [0, NOPE_DIM) is a deterministic per-tensor fp8 quant

716 + # of the (un-rotated) KV input, so it must be bit-identical. The RoPE region

717 + # [NOPE_DIM, HEAD_DIM) is rotated in fp32 and may differ by <=1 fp8 ULP.

718 + k_fused = _as_stored_fp8(k_cache_fused)

719 torch.testing.assert_close(

603 - q_fp8_fused.float(), q_fp8_ref.float(), rtol=0, atol=0.25

720 + k_fused[..., :NOPE_DIM].float(),

721 + k_cache_ref[..., :NOPE_DIM].float(),

722 + rtol=0,

723 + atol=0,

724 )

605 - torch.testing.assert_close(

606 - k_cache_fused.float(), k_cache_ref.float(), rtol=0, atol=0.25

725 + k_max_ulp = int(

726 + _fp8_ulp_distance(k_fused[..., NOPE_DIM:], k_cache_ref[..., NOPE_DIM:])

727 + .max()

728 + .item()

729 )

730 + assert k_max_ulp <= 1, f"K-cache RoPE fp8 differs by {k_max_ulp} ULP (>1)"

731

732

733 @pytest.mark.skipif(

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+22 -3

@@ -1363,9 +1363,28 @@ def process_fp8_weight_block_strategy(

1363 )

1364

1365 if current_platform.is_fp8_fnuz() and weight.dtype == torch.float8_e4m3fn:

1366 - weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(

1367 - weight=weight, weight_scale=weight_scale

1368 - )

1366 + if weight_scale.dtype == torch.float8_e8m0fnu:

1367 + # UE8M0 scales: e8m0 stores exponent-only values (2^(exp-127)),

1368 + # so doubling the dequant scale == incrementing the exponent byte

1369 + # by 1. Convert the OCP E4M3 weight bytes to FNUZ in place by

1370 + # reinterpreting and patching the NaN sentinel (-128 in int8),

1371 + # then double the UE8M0 exponent so the dequantized magnitudes

1372 + # match.

1373 + weight_as_int8 = weight.view(torch.int8)

1374 + ROCM_FP8_NAN_AS_INT = -128

1375 + weight_as_int8[weight_as_int8 == ROCM_FP8_NAN_AS_INT] = 0

1376 + weight = weight_as_int8.view(torch.float8_e4m3fnuz)

1377 + exp_bytes = weight_scale.view(torch.uint8)

1378 + weight_scale = (

1379 + (exp_bytes.to(torch.int16) + 1)

1380 + .clamp(max=254)

1381 + .to(torch.uint8)

1382 + .view(torch.float8_e8m0fnu)

1383 + )

1384 + else:

1385 + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(

1386 + weight=weight, weight_scale=weight_scale

1387 + )

1388

1389 weight = _maybe_pad_fp8_weight(weight)

1390 return weight, weight_scale

vllm/models/deepseek_v4/amd/rocm.py

@@ -14,6 +14,7 @@ from vllm.models.deepseek_v4.sparse_mla import (

14 DeepseekV4FlashMLAMetadata,

15 DeepseekV4FlashMLAMetadataBuilder,

16 )

17 +from vllm.platforms import current_platform

18 from vllm.triton_utils import tl, triton

19 from vllm.v1.attention.backend import (

20 CommonAttentionMetadata,

@@ -796,6 +797,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

797 assert attn_metadata is not None

798 assert compressed_k_cache is not None

799 block_table = attn_metadata.block_table[num_decodes:]

800 + # compressed_k_cache is OCP on every platform (Triton encoder).

801 dequantize_and_gather_k_cache(

802 kv[:chunk_size],

803 compressed_k_cache,

@@ -804,6 +806,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

806 block_table=block_table[chunk_start:chunk_end],

807 block_size=attn_metadata.block_size // self.compress_ratio,

808 offset=0,

809 + use_fnuz=False,

810 )

811

812 swa_block_table = swa_metadata.block_table[num_decodes:]

@@ -815,6 +818,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

818 block_table=swa_block_table[chunk_start:chunk_end],

819 block_size=swa_metadata.block_size,

820 offset=N,

821 + use_fnuz=current_platform.is_fp8_fnuz(),

822 )

823

824 query_start = (

vllm/models/deepseek_v4/common/ops/cache_utils.py

+49 -6

@@ -16,6 +16,10 @@ preparation.

17 import torch

19 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

20 + get_fp8_min_max,

21 +)

22 +from vllm.platforms import current_platform

23 from vllm.triton_utils import tl, triton

24 from vllm.utils.import_utils import has_cutedsl

@@ -39,6 +43,7 @@ def quantize_and_insert_k_kernel(

43 block_stride: tl.constexpr, # total bytes per block (padded)

44 fp8_max: tl.constexpr,

45 n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding)

46 + use_fnuz: tl.constexpr = False,

47 ):

48 """

49 Quantize K tensor and insert into paged K cache.

@@ -49,6 +54,9 @@ def quantize_and_insert_k_kernel(

54 - [64*576 + 64*8, block_stride): Padding

56 One program per token.

57 +

58 + ``use_fnuz=True`` selects FNUZ (``tl.float8e4b8``); default OCP

59 + (``tl.float8e4nv``) matches every production caller.

60 """

61 pid = tl.program_id(0)

@@ -112,8 +120,11 @@ def quantize_and_insert_k_kernel(

120 x_scaled = x / scale

121 x_clamped = tl.clamp(x_scaled, -fp8_max, fp8_max)

122

115 - # Convert to fp8, then bitcast to uint8 for storage

116 - x_fp8 = x_clamped.to(tl.float8e4nv)

123 + # Convert to fp8 (FNUZ on gfx942, OCP elsewhere), then bitcast to uint8.

124 + if use_fnuz:

125 + x_fp8 = x_clamped.to(tl.float8e4b8)

126 + else:

127 + x_fp8 = x_clamped.to(tl.float8e4nv)

128 x_uint8 = x_fp8.to(tl.uint8, bitcast=True)

129

130 # Store as uint8 (1 byte each)

@@ -145,6 +156,7 @@ def quantize_and_insert_k_cache(

156 slot_mapping: torch.Tensor, # [num_tokens] int64

157 block_size: int = 64,

158 is_ue8m0: bool = True,

159 + use_fnuz: bool = False,

160 ):

161 """

162 Quantize K tensor and insert into paged K cache.

@@ -155,6 +167,10 @@ def quantize_and_insert_k_cache(

167 - Next 64 * 8 = 512 bytes: Scales

168 - Each token: 8 bytes (uint8 scales, 7 real + 1 padding)

169 - Padded to multiple of 576

170 +

171 + ``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on

172 + platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3,

173 + which is used by OCP-encoded caches even on gfx942.

174 """

175 assert k.dim() == 2 and k.shape[1] == 512, (

176 f"K must be [num_tokens, 512], got {k.shape}"

@@ -171,7 +187,12 @@ def quantize_and_insert_k_cache(

187 TOKEN_BF16_DIM = 64

188 TOKEN_SCALE_DIM = 8

189 QUANT_BLOCK_SIZE = 64

174 - FP8_MAX = 448.0

190 + if use_fnuz:

191 + if not current_platform.is_fp8_fnuz():

192 + raise ValueError("use_fnuz=True requires a platform using FNUZ FP8")

193 + _, FP8_MAX = get_fp8_min_max()

194 + else:

195 + FP8_MAX = torch.finfo(torch.float8_e4m3fn).max

196 TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2

197

198 grid = (num_tokens,)

@@ -191,6 +212,7 @@ def quantize_and_insert_k_cache(

212 block_stride=block_stride,

213 fp8_max=FP8_MAX,

214 n_quant_blocks=8,

215 + use_fnuz=use_fnuz,

216 )

217

218

@@ -216,6 +238,7 @@ def _dequantize_and_gather_k_kernel(

238 output_dim: tl.constexpr, # 512

239 fp8_max: tl.constexpr,

240 n_quant_blocks: tl.constexpr, # 7 real blocks

241 + use_fnuz: tl.constexpr = False,

242 ):

243 batch_idx = tl.program_id(0)

244 worker_id = tl.program_id(1)

@@ -273,8 +296,11 @@ def _dequantize_and_gather_k_kernel(

296 # Load quantized fp8 values (stored as uint8)

297 x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0)

298

276 - # Bitcast uint8 back to fp8

277 - x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)

299 + # Bitcast uint8 back to fp8 (FNUZ on gfx942, OCP elsewhere).

300 + if use_fnuz:

301 + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)

302 + else:

303 + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)

304

305 # Convert fp8 to float32 for computation

306 x_float = x_fp8.to(tl.float32)

@@ -317,6 +343,7 @@ def dequantize_and_gather_k_cache_triton(

343 block_table: torch.Tensor,

344 block_size: int,

345 offset: int,

346 + use_fnuz: bool = False,

347 ) -> None:

348 TOKEN_FP8_DIM = 448

349 TOKEN_BF16_DIM = 64

@@ -347,6 +374,7 @@ def dequantize_and_gather_k_cache_triton(

374 output_dim=512,

375 fp8_max=FP8_MAX,

376 n_quant_blocks=7,

377 + use_fnuz=use_fnuz,

378 )

379

380

@@ -363,7 +391,15 @@ def dequantize_and_gather_k_cache(

391 block_table: torch.Tensor,

392 block_size: int,

393 offset: int,

394 + use_fnuz: bool = False,

395 ) -> None:

396 + """Dequantize and gather a paged DSv4 K cache.

397 +

398 + ``use_fnuz`` MUST match the encoder of the specific cache being read:

399 + ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere),

400 + ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder

401 + writes FNUZ on gfx942 and OCP on gfx950).

402 + """

403 if has_cutedsl():

404 # lazily import, otherwise some tests fail due to CUDA driver init failure.

405 from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (

@@ -376,7 +412,14 @@ def dequantize_and_gather_k_cache(

412 return

413

414 dequantize_and_gather_k_cache_triton(

379 - out, k_cache, seq_lens, gather_lens, block_table, block_size, offset

415 + out,

416 + k_cache,

417 + seq_lens,

418 + gather_lens,

419 + block_table,

420 + block_size,

421 + offset,

422 + use_fnuz=use_fnuz,

423 )

424

425

vllm/models/deepseek_v4/nvidia/ops/o_proj.py

+3 -1

@@ -3,7 +3,9 @@

3 import torch

4 import torch.nn as nn

6 -from vllm.models.deepseek_v4.common.ops import fused_inv_rope_fp8_quant

6 +from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (

7 + fused_inv_rope_fp8_quant,

8 +)

9 from vllm.platforms import current_platform

10 from vllm.utils.deep_gemm import fp8_einsum

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

+47 -14

@@ -504,7 +504,13 @@ def fp8_mqa_logits_torch(

504 )

505 mask = mask_lo & mask_hi

506

507 - score = torch.einsum("mhd,nd->hmn", q, k).float() * scale

507 + # ``score`` is [H, M, N]; ``scale`` is the per-KV-token scale, which

508 + # vLLM callers hand us as ``[N, 1]`` (a ``[N, 4]`` uint8 buffer cast

509 + # to fp32). PyTorch right-aligns dimensions for broadcasting, so a

510 + # naked ``score * scale`` would align ``scale``'s leading dim with

511 + # ``score``'s M dim and raise a shape mismatch. Flatten to ``[N]`` so

512 + # broadcasting lines up with the last dim of ``score``.

513 + score = torch.einsum("mhd,nd->hmn", q, k).float() * scale.reshape(-1)

514 logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)

515 logits = logits.masked_fill(~mask, float("-inf"))

516

@@ -557,13 +563,26 @@ def rocm_fp8_mqa_logits(

563 # path after aiter merge this kernel into main

564 from vllm._aiter_ops import rocm_aiter_ops

565

566 + k_fp8, scale = kv

567 +

568 + # Temporarily route gfx942 to the vendored ROCm/aiter#3257 workaround.

569 + # Remove this branch once vLLM bumps AITER to a version that includes

570 + # ROCm/aiter#3257.

571 + if _ON_GFX942 and rocm_aiter_ops.is_enabled():

572 + from vllm.v1.attention.ops.triton_fp8_mqa_logits import (

573 + fp8_mqa_logits_gfx942,

574 + )

575 +

576 + return fp8_mqa_logits_gfx942(

577 + q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke

578 + )

579 +

580 aiter_mqa_logits_module = None

581 if rocm_aiter_ops.is_enabled():

582 aiter_mqa_logits_module = mqa_logits_module()

583

584 if aiter_mqa_logits_module is not None:

585 fp8_mqa_logits = aiter_mqa_logits_module.fp8_mqa_logits

566 - k_fp8, scale = kv

586 return fp8_mqa_logits(q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke)

587 else:

588 return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)

@@ -1249,7 +1268,10 @@ def _sparse_attn_decode_ragged_kernel(

1268 NOPE_DIM: tl.constexpr,

1269 NOPE_BLOCK: tl.constexpr,

1270 ROPE_DIM: tl.constexpr,

1252 - IS_FNUZ: tl.constexpr,

1271 + # SWA K-cache (main): C++ encoder writes FNUZ on gfx942, OCP on gfx950.

1272 + # Compressed K-cache (extra): Triton encoder writes OCP everywhere.

1273 + IS_FNUZ_MAIN: tl.constexpr,

1274 + IS_FNUZ_EXTRA: tl.constexpr,

1275 BLOCK_H: tl.constexpr,

1276 BLOCK_K: tl.constexpr,

1277 ):

@@ -1306,8 +1328,8 @@ def _sparse_attn_decode_ragged_kernel(

1328 mask=valid[:, None] & nope_mask[None, :],

1329 other=0,

1330 )

1309 - if IS_FNUZ:

1310 - x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True)

1331 + if IS_FNUZ_MAIN:

1332 + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)

1333 else:

1334 x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)

1335 encoded_scales = tl.load(

@@ -1374,8 +1396,8 @@ def _sparse_attn_decode_ragged_kernel(

1396 mask=valid[:, None] & nope_mask[None, :],

1397 other=0,

1398 )

1377 - if IS_FNUZ:

1378 - x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True)

1399 + if IS_FNUZ_EXTRA:

1400 + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)

1401 else:

1402 x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)

1403 encoded_scales = tl.load(

@@ -1485,7 +1507,12 @@ def _sparse_attn_decode_partial_kernel(

1507 NOPE_DIM: tl.constexpr,

1508 NOPE_BLOCK: tl.constexpr,

1509 ROPE_DIM: tl.constexpr,

1488 - IS_FNUZ: tl.constexpr,

1510 + # `main_cache` is the SWA K-cache (written by the C++ encoder, FNUZ on

1511 + # gfx942 / OCP on gfx950). `extra_cache` is the compressed K-cache

1512 + # (Triton encoder, OCP on every platform). Reading both with the same

1513 + # `IS_FNUZ` would decode one of them with the wrong FNUZ/OCP scale ratio.

1514 + IS_FNUZ_MAIN: tl.constexpr,

1515 + IS_FNUZ_EXTRA: tl.constexpr,

1516 BLOCK_H: tl.constexpr,

1517 BLOCK_K: tl.constexpr,

1518 NUM_SPLITS: tl.constexpr,

@@ -1551,8 +1578,8 @@ def _sparse_attn_decode_partial_kernel(

1578 mask=valid[:, None] & nope_mask[None, :],

1579 other=0,

1580 )

1554 - if IS_FNUZ:

1555 - x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True)

1581 + if IS_FNUZ_MAIN:

1582 + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)

1583 else:

1584 x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)

1585 encoded_scales = tl.load(

@@ -1622,8 +1649,8 @@ def _sparse_attn_decode_partial_kernel(

1649 mask=valid[:, None] & nope_mask[None, :],

1650 other=0,

1651 )

1625 - if IS_FNUZ:

1626 - x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True)

1652 + if IS_FNUZ_EXTRA:

1653 + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)

1654 else:

1655 x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)

1656 encoded_scales = tl.load(

@@ -2095,7 +2122,8 @@ def _rocm_sparse_attn_decode_ragged_triton(

2122 NOPE_DIM=nope_head_dim,

2123 NOPE_BLOCK=nope_block,

2124 ROPE_DIM=rope_head_dim,

2098 - IS_FNUZ=is_fnuz,

2125 + IS_FNUZ_MAIN=is_fnuz,

2126 + IS_FNUZ_EXTRA=False,

2127 BLOCK_H=block_h,

2128 BLOCK_K=block_k,

2129 num_warps=8,

@@ -2153,7 +2181,12 @@ def _rocm_sparse_attn_decode_ragged_triton(

2181 NOPE_DIM=nope_head_dim,

2182 NOPE_BLOCK=nope_block,

2183 ROPE_DIM=rope_head_dim,

2156 - IS_FNUZ=is_fnuz,

2184 + # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950).

2185 + # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere).

2186 + # Reading both with a single IS_FNUZ would decode one of them with the

2187 + # wrong FNUZ/OCP scale ratio (~1.87×).

2188 + IS_FNUZ_MAIN=is_fnuz,

2189 + IS_FNUZ_EXTRA=False,

2190 BLOCK_H=block_h,

2191 BLOCK_K=block_k,

2192 NUM_SPLITS=num_splits,

vllm/v1/attention/ops/triton_fp8_mqa_logits.py new

+262

@@ -0,0 +1,262 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Temporary gfx942 fallback for AITER's fp8_mqa_logits kernel.

4 +

5 +This module vendors AITER's Triton fp8_mqa_logits kernel with the gfx942

6 +tile-size workaround from ROCm/aiter#3257. It is used only while vLLM's

7 +pinned AITER version lacks that fix.

8 +

9 +TODO: Remove this vendored copy once vLLM pins an AITER version that includes

10 +ROCm/aiter#3257 bugfix for gfx942.

11 +"""

12 +

13 +import torch

14 +

15 +from vllm.triton_utils import tl, triton

16 +

17 +# gfx942 (MI300X) has 64 KiB of LDS per CU. We accept the default

18 +# (BLOCK_KV=128, num_stages=2) tile only when *both* of these hold:

19 +#

20 +# 1. Occupancy gate. With waves_per_eu=2 and num_warps=4 we target two

21 +# workgroups co-resident on a CU -> per-WG LDS budget = 32 KiB. Triton

22 +# keeps Q in registers (loop-invariant) and the fp32 scores accumulator

23 +# in VGPRs (heavy VALU), so only the double-buffered KV tile is

24 +# expected to live in LDS. A 0.9 safety factor leaves headroom for any

25 +# LDS overhead the compiler may add.

26 +#

27 +# 2. Hardware ceiling. Defensive upper bound that also counts Q and

28 +# scores against the 64 KiB CU limit, in case a Triton version (older

29 +# or future) decides to spill them to LDS. False positives here only

30 +# shrink the tile; false negatives are JIT-aborts, so we lean

31 +# conservative.

32 +_GFX942_CU_LDS_BYTES = 64 * 1024

33 +_GFX942_PER_WG_LDS_BUDGET_BYTES = _GFX942_CU_LDS_BYTES * 9 // 20 # ~28.8 KiB

34 +

35 +

36 +def _gfx942_default_tile_fits_lds(num_heads: int, head_size: int) -> bool:

37 + """Return True iff (BLOCK_KV=128, num_stages=2) fits in MI300X LDS."""

38 + BLOCK_KV = 128

39 + NUM_STAGES = 2

40 + kv_bytes = head_size * BLOCK_KV * NUM_STAGES

41 + scores_bytes = num_heads * BLOCK_KV * 4

42 + q_bytes = num_heads * head_size

43 + fits_occupancy = kv_bytes < _GFX942_PER_WG_LDS_BUDGET_BYTES

44 + fits_hardware = q_bytes + kv_bytes + scores_bytes <= _GFX942_CU_LDS_BYTES

45 + return fits_occupancy and fits_hardware

46 +

47 +

48 +@triton.jit

49 +def _fp8_mqa_logits_kernel(

50 + Q_ptr, # fp8e4m3 [seq_len, H, D]

51 + KV_ptr, # fp8e4m3 [seq_len_kv, D]

52 + kv_scales_ptr, # fp32 [seq_len_kv]

53 + weights_ptr, # fp32 [seq_len, H]

54 + cu_start_ptr, # int32 [seq_len]

55 + cu_end_ptr, # int32 [seq_len]

56 + logits_ptr, # fp32 [seq_len, seq_len_kv]

57 + seq_len,

58 + seq_len_kv,

59 + NUM_HEADS: tl.constexpr,

60 + HEAD_SIZE: tl.constexpr,

61 + # strides

62 + stride_q_s: tl.int64,

63 + stride_q_h: tl.constexpr,

64 + stride_q_d: tl.constexpr,

65 + stride_kv_s: tl.int64,

66 + stride_kv_d: tl.constexpr,

67 + stride_w_s: tl.int64,

68 + stride_w_h: tl.constexpr,

69 + stride_logits_s: tl.int64,

70 + stride_logits_k: tl.int64,

71 + # block sizes

72 + BLOCK_KV: tl.constexpr,

73 +):

74 + row_id = tl.program_id(0)

75 + # go from larger to smaller in terms of work

76 + # to reduce the tail effect

77 + row_id = tl.num_programs(0) - row_id - 1

78 + tl.assume(row_id >= 0)

79 + tl.assume(stride_q_s > 0)

80 + tl.assume(stride_q_h > 0)

81 + tl.assume(stride_q_d > 0)

82 + tl.assume(stride_kv_s > 0)

83 + tl.assume(stride_kv_d > 0)

84 + tl.assume(stride_w_s > 0)

85 + tl.assume(stride_w_h > 0)

86 +

87 + logits_row_ptrs = logits_ptr + row_id * stride_logits_s

88 +

89 + h_inds = tl.arange(0, NUM_HEADS)[:, None]

90 + d_inds = tl.arange(0, HEAD_SIZE)

91 +

92 + # load Q[BLOCK_Q, NUM_HEADS, HEAD_SIZE]

93 + q_ptrs = (

94 + Q_ptr + row_id * stride_q_s + h_inds * stride_q_h + d_inds[None, :] * stride_q_d

95 + )

96 +

97 + q_block = tl.load(q_ptrs, cache_modifier=".cg")

98 + w_ptrs = weights_ptr + row_id * stride_w_s + h_inds * stride_w_h

99 + w_block = tl.load(w_ptrs, cache_modifier=".cg").to(tl.float32)

100 +

101 + # Load start/end for each row in this block

102 + start_ind = tl.load(cu_start_ptr + row_id)

103 + end_ind = tl.load(cu_end_ptr + row_id)

104 +

105 + start_ind = tl.maximum(start_ind, 0)

106 + end_ind = tl.minimum(end_ind, seq_len_kv)

107 + shifted_end = end_ind - start_ind

108 + shifted_unmasked_end = shifted_end // BLOCK_KV * BLOCK_KV

109 +

110 + kv_col_offsets = tl.arange(0, BLOCK_KV) + start_ind

111 + kv_ptrs = (

112 + KV_ptr + kv_col_offsets[None, :] * stride_kv_s + d_inds[:, None] * stride_kv_d

113 + )

114 +

115 + kv_scales_ptrs = kv_scales_ptr + kv_col_offsets

116 +

117 + logits_ptrs = logits_row_ptrs + kv_col_offsets * stride_logits_k

118 +

119 + # Loop over KV tiles

120 + for _ in tl.range(0, shifted_unmasked_end, BLOCK_KV):

121 + kv_block = tl.load(kv_ptrs)

122 + kv_scales = tl.load(kv_scales_ptrs)

123 +

124 + # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]

125 + scores = tl.dot(q_block, kv_block, input_precision="ieee")

126 + # Multiply by kv_scales (broadcast along rows)

127 + scores = scores * kv_scales[None, :]

128 + # ReLU

129 + scores = tl.maximum(scores, 0.0)

130 + scores = scores * w_block

131 + # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ]

132 + scores = tl.sum(scores, axis=0)

133 + tl.store(logits_ptrs, scores)

134 +

135 + kv_ptrs += BLOCK_KV * stride_kv_s

136 + kv_scales_ptrs += BLOCK_KV

137 + logits_ptrs += BLOCK_KV * stride_logits_k

138 + kv_col_offsets += BLOCK_KV

139 +

140 + # masked load

141 + kv_col_mask = kv_col_offsets < end_ind

142 + kv_block = tl.load(kv_ptrs, mask=kv_col_mask[None, :], other=0.0)

143 + kv_scales = tl.load(kv_scales_ptrs, mask=kv_col_mask, other=0.0)

144 +

145 + # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]

146 + scores = tl.dot(q_block, kv_block, input_precision="ieee")

147 + # Multiply by kv_scales (broadcast along rows)

148 + scores = scores * kv_scales[None, :]

149 + # ReLU

150 + scores = tl.maximum(scores, 0.0)

151 + scores = scores * w_block

152 + # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ]

153 + scores = tl.sum(scores, axis=0)

154 + # masked store

155 + in_window = (kv_col_offsets >= start_ind) & (kv_col_offsets < end_ind)

156 + tl.store(logits_ptrs, scores, mask=in_window)

157 +

158 +

159 +def fp8_mqa_logits_gfx942(

160 + q: torch.Tensor,

161 + k_fp8: torch.Tensor,

162 + kv_scales: torch.Tensor,

163 + weights: torch.Tensor,

164 + cu_starts: torch.Tensor,

165 + cu_ends: torch.Tensor,

166 +) -> torch.Tensor:

167 + """Compute FP8 MQA logits on MI300X (gfx942) using the vendored kernel.

168 +

169 + Drop-in replacement for ``aiter.ops.triton.attention.fp8_mqa_logits.

170 + fp8_mqa_logits`` on MI300X. Selects ``(BLOCK_KV, num_stages)`` based on

171 + whether the default tile fits within the 64 KiB LDS budget of a gfx942

172 + CU (see module docstring).

173 +

174 + Args:

175 + q: Query tensor of shape ``[M, H, D]``, FP8 dtype.

176 + k_fp8: Key tensor of shape ``[N, D]``, FP8 dtype.

177 + kv_scales: K scales of shape ``[N]`` (or ``[N, 1]`` -- viewed as

178 + ``[N]``), float32.

179 + weights: Per-head weights of shape ``[M, H]``, float32.

180 + cu_starts: Start indices (inclusive) of shape ``[M]``, int32.

181 + cu_ends: End indices (exclusive) of shape ``[M]``, int32.

182 +

183 + Returns:

184 + Logits of shape ``[M, N]``, float32 -- positions outside

185 + ``[cu_starts[i], cu_ends[i])`` for row ``i`` are pre-filled with

186 + ``-inf`` so the caller can run a top-k without masking.

187 + """

188 + seq_len, num_heads, head_size = q.shape

189 + seq_len_kv = k_fp8.shape[0]

190 + assert num_heads & (num_heads - 1) == 0, (

191 + f"num_heads must be a power of two (got {num_heads})"

192 + )

193 + assert head_size & (head_size - 1) == 0, (

194 + f"head_size must be a power of two (got {head_size})"

195 + )

196 +

197 + # The kernel walks ``kv_scales`` as a 1-D contiguous array of size N

198 + # (it indexes by ``kv_scales_ptr + kv_col_offsets``). The vLLM caller

199 + # passes a ``[N, 4]`` uint8 view-cast-to-float32 which lands as

200 + # ``[N, 1]`` contiguous -- byte-identical to ``[N]`` -- but flatten

201 + # explicitly to keep the kernel's pointer arithmetic intent clear.

202 + kv_scales_1d = kv_scales.reshape(-1)

203 +

204 + # Initialise with -inf so positions outside [cu_starts, cu_ends) read

205 + # as ``-inf`` after the masked store path -- this matches AITER's

206 + # ``fp8_mqa_logits`` semantics and is what the top-k consumer expects.

207 + logits = torch.full(

208 + (seq_len, seq_len_kv),

209 + fill_value=-float("inf"),

210 + dtype=torch.float32,

211 + device=q.device,

212 + )

213 +

214 + if _gfx942_default_tile_fits_lds(num_heads, head_size):

215 + block_kv = 128

216 + num_stages = 2

217 + else:

218 + # DSv4 sparse indexer (NUM_HEADS=64, HEAD_SIZE=128) lands here:

219 + # default tile spills past gfx942's 64 KiB LDS budget. (64, 1)

220 + # needs ~33 KiB and clears the per-WG budget with margin.

221 + block_kv = 64

222 + num_stages = 1

223 +

224 + # heuristic for MFMA instruction shape, identical to AITER's choice

225 + matrix_instr_nonkdim = 32

226 + if seq_len <= 1024:

227 + matrix_instr_nonkdim = 16

228 +

229 + stride_q_s, stride_q_h, stride_q_d = q.stride()

230 + stride_kv_s, stride_kv_d = k_fp8.stride()

231 + stride_w_s, stride_w_h = weights.stride()

232 + stride_logits_s, stride_logits_k = logits.stride()

233 +

234 + _fp8_mqa_logits_kernel[(seq_len,)](

235 + Q_ptr=q,

236 + KV_ptr=k_fp8,

237 + kv_scales_ptr=kv_scales_1d,

238 + weights_ptr=weights,

239 + cu_start_ptr=cu_starts,

240 + cu_end_ptr=cu_ends,

241 + logits_ptr=logits,

242 + seq_len=seq_len,

243 + seq_len_kv=seq_len_kv,

244 + NUM_HEADS=num_heads,

245 + HEAD_SIZE=head_size,

246 + stride_q_s=stride_q_s,

247 + stride_q_h=stride_q_h,

248 + stride_q_d=stride_q_d,

249 + stride_kv_s=stride_kv_s,

250 + stride_kv_d=stride_kv_d,

251 + stride_w_s=stride_w_s,

252 + stride_w_h=stride_w_h,

253 + stride_logits_s=stride_logits_s,

254 + stride_logits_k=stride_logits_k,

255 + BLOCK_KV=block_kv,

256 + num_warps=4,

257 + num_stages=num_stages,

258 + waves_per_eu=2,

259 + matrix_instr_nonkdim=matrix_instr_nonkdim,

260 + )

261 +

262 + return logits