2267d3b112 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -21,6 +21,7 @@ from vllm.model_executor.layers.hpc.hpc_module import HpcModule

21 from vllm.utils.torch_utils import direct_register_custom_op

22 from vllm.v1.attention.backends.hpc_attn import HpcAttnMetadata

23 from vllm.v1.attention.backends.registry import AttentionBackendEnum

24 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

25

26 logger = init_logger(__name__)

27

@@ -280,6 +281,111 @@ class HpcRopeNorm(CustomOp, HpcModule):

281 torch.ops.vllm.hpc_rope_norm_forward(qkv, output, layer_name)

282 return output

283

284 + def _kv_write_scratch(self, cache: torch.Tensor) -> torch.Tensor:

285 + """Throwaway dense K/V pages for the stride-aware fallback path.

286 +

287 + ``hpc.rope_norm_store_kv[_fp8]`` always writes the paged cache itself

288 + (rotated K/V *and* zero padding for the tail of every request's last

289 + block), addressing it densely. When the real pages are not dense those

290 + writes would land on unrelated blocks, so they are redirected here by

291 + pairing this scratch cache with an all-zero block table; the real data

292 + is taken out through ``out_k``/``out_v`` instead.

293 +

294 + Returns a (2, 1, block_size, num_kv_heads, head_dim) tensor: index 0 is

295 + the key scratch page, index 1 the value scratch page.

296 + """

297 + block_size = cache.shape[1]

298 + key = (cache.dtype, block_size, cache.device)

299 + if getattr(self, "_kv_scratch_key", None) != key:

300 + self._kv_scratch = torch.zeros(

301 + 2,

302 + 1,

303 + block_size,

304 + self.num_kv_heads,

305 + self.head_dim,

306 + dtype=cache.dtype,

307 + device=cache.device,

308 + )

309 + self._kv_scratch_key = key

310 + return self._kv_scratch

311 +

312 + @staticmethod

313 + def _zero_pad_last_blocks(

314 + key_cache: torch.Tensor,

315 + value_cache: torch.Tensor,

316 + seq_lens: torch.Tensor,

317 + block_table: torch.Tensor,

318 + ) -> None:

319 + """Zero the unused tail slots of each request's last block.

320 +

321 + The HPC attention kernels read whole blocks, so the slots past

322 + ``seq_len`` must be zero. The fused op normally does this itself; on the

323 + fallback path its (densely addressed) padding writes are discarded, so

324 + the padding is reapplied here through the real, strided cache views.

325 +

326 + Shapes are static and no value is inspected on the host, so this stays

327 + CUDA-graph capturable. It is also self-contained: each request's last

328 + block is read, masked and written back, so slots that must be preserved

329 + (and requests padded in for graph capture, whose ``seq_len`` is 0) keep

330 + their previous contents instead of being redirected somewhere.

331 + """

332 + block_size = key_cache.shape[1]

333 + seq_lens = seq_lens.long()

334 + # seq_len == 0 only happens for graph-capture padding; clamping keeps

335 + # the gather in bounds and the all-False mask below makes the write a

336 + # no-op for those rows.

337 + last_block_idx = (seq_lens - 1).clamp_min(0) // block_size

338 + last_block_idx = last_block_idx.clamp_max(block_table.shape[1] - 1)

339 + last_block = block_table.gather(1, last_block_idx.unsqueeze(1)).squeeze(1)

340 + last_block = last_block.long()

341 +

342 + remainder = seq_lens % block_size

343 + # A block that ends exactly on the boundary has no tail to clear.

344 + remainder = torch.where(remainder == 0, block_size, remainder)

345 +

346 + offsets = torch.arange(block_size, device=key_cache.device)

347 + is_tail = (offsets.unsqueeze(0) >= remainder.unsqueeze(1)) & (

348 + seq_lens > 0

349 + ).unsqueeze(1)

350 + is_tail = is_tail[..., None, None]

351 +

352 + for cache in (key_cache, value_cache):

353 + view = cache

354 + if view.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):

355 + view = view.view(torch.uint8)

356 + blocks = view.index_select(0, last_block)

357 + blocks = torch.where(is_tail, blocks.new_zeros(()), blocks)

358 + view.index_copy_(0, last_block, blocks)

359 +

360 + @staticmethod

361 + def _scatter_kv_cache(

362 + cache: torch.Tensor,

363 + src: torch.Tensor,

364 + slot_mapping: torch.Tensor,

365 + ) -> None:

366 + """Write dense per-token K (or V) into a paged cache view.

367 +

368 + Used on the stride-aware fallback path (see ``_forward_impl``).

369 + ``cache`` is a (num_blocks, block_size, num_kv_heads, head_dim) view

370 + with arbitrary block/token/head strides; ``src`` is the dense

371 + (num_tokens, num_kv_heads, head_dim) buffer produced by the fused HPC

372 + kernel. FP8 payloads are moved as raw bytes because ``index_copy_``

373 + has no FP8 kernel.

374 +

375 + Padded tokens carry ``PAD_SLOT_ID`` (-1); clamping sends them to the

376 + reserved null block instead of tripping ``index_copy_``'s bounds check.

377 + This keeps the op CUDA-graph safe (static shape, no host sync), unlike

378 + masking the padded rows out.

379 + """

380 + # (num_blocks, block_size, ...) -> (num_slots, num_kv_heads, head_dim).

381 + # Valid as a view for any layout whose block stride is a whole number

382 + # of token strides, which holds for every layout the backend accepts.

383 + flat = cache.flatten(0, 1)

384 + if flat.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):

385 + flat = flat.view(torch.uint8)

386 + src = src.view(torch.uint8)

387 + flat.index_copy_(0, slot_mapping.clamp_min(NULL_BLOCK_ID), src)

388 +

389 def _forward_impl(

390 self,

391 qkv: torch.Tensor,

@@ -295,9 +401,32 @@ class HpcRopeNorm(CustomOp, HpcModule):

401 """

402 import hpc

403

404 + # KV cache for the FP8 path is stored as uint8; view it as fp8 so the

405 + # rope_norm_store_kv_fp8 kernel can write quantized K/V in-place.

406 + # Must happen *before* deriving the K/V views, so that they (and any

407 + # buffer whose dtype is taken from them) carry the FP8 dtype.

408 + if self.use_fp8:

409 + kv_cache = kv_cache.view(torch.float8_e4m3fn)

410 +

411 # (B, H, N, 2*hs) -> two (B, N, H, hs) K/V views, as in HpcAttentionImpl.

412 key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_dim, dim=-1)

413

414 + # hpc.rope_norm_store_kv[_fp8] addresses its K/V caches as *dense*

415 + # NHD tensors and ignores their strides, so it can only write the cache

416 + # in place when both views are contiguous. Packed layouts (the backend's

417 + # LBNHC pages interleave K and V inside the state-content dim, so each

418 + # split view has a token stride of 2 * head_dim) must instead take the

419 + # rotated K/V out through out_k/out_v and be scattered separately.

420 + fused_kv_write = key_cache.is_contiguous() and value_cache.is_contiguous()

421 + if not fused_kv_write:

422 + logger.warning_once(

423 + "HpcRopeNorm: KV cache pages interleave K and V, which "

424 + "hpc.rope_norm_store_kv[_fp8] cannot address, so K/V are "

425 + "written by a separate scatter. Update hpc-ops to a build with "

426 + "strided KV cache support to re-enable the fully fused write."

427 + )

428 + slot_mapping = attn_metadata.slot_mapping

429 +

430 num_actual_tokens = attn_metadata.num_actual_tokens

431 num_prefill_reqs = attn_metadata.num_prefills

432 num_decode_reqs = attn_metadata.num_decodes

@@ -307,11 +436,6 @@ class HpcRopeNorm(CustomOp, HpcModule):

436

437 num_prefill_tokens = num_actual_tokens - num_decode_tokens

438

310 - # KV cache for the FP8 path is stored as uint8; view it as fp8 so the

311 - # rope_norm_store_kv_fp8 kernel can write quantized K/V in-place.

312 - if self.use_fp8:

313 - kv_cache = kv_cache.view(torch.float8_e4m3fn)

314 -

439 # Per-tensor K/V scales (shape [1]) used by the FP8 kernel.

440 k_scale = attn_layer._k_scale.reshape(1)

441 v_scale = attn_layer._v_scale.reshape(1)

@@ -334,15 +458,35 @@ class HpcRopeNorm(CustomOp, HpcModule):

458 num_decode_tokens : num_decode_tokens + num_prefill_tokens

459 ]

460

461 + kv_dtype = key_cache.dtype

462 + if fused_kv_write:

463 + out_k = out_v = None

464 + kcache_arg, vcache_arg = key_cache, value_cache

465 + kv_indices_arg = block_table_prefill

466 + else:

467 + out_k = torch.empty(

468 + num_prefill_tokens,

469 + self.num_kv_heads,

470 + self.head_dim,

471 + dtype=kv_dtype,

472 + device=qkv.device,

473 + )

474 + out_v = torch.empty_like(out_k)

475 + # Send the kernel's own cache writes to scratch (all block ids

476 + # remapped to the single scratch page).

477 + scratch = self._kv_write_scratch(key_cache)

478 + kcache_arg, vcache_arg = scratch[0], scratch[1]

479 + kv_indices_arg = torch.zeros_like(block_table_prefill)

480 +

481 if self.use_fp8:

482 _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(

339 - key_cache=key_cache,

340 - value_cache=value_cache,

483 + key_cache=kcache_arg,

484 + value_cache=vcache_arg,

485 qkv=qkv_prefill,

486 cos_sin=self.cos_sin_cache,

487 num_seqlen_per_req=seq_lens_prefill,

488 q_index=cu_seqlens_prefill,

345 - kvcache_indices=block_table_prefill,

489 + kvcache_indices=kv_indices_arg,

490 is_prefill=True,

491 k_scale=k_scale,

492 v_scale=v_scale,

@@ -352,24 +496,36 @@ class HpcRopeNorm(CustomOp, HpcModule):

496 k_norm_weight=k_norm_weight,

497 qk_norm_policy=self.qk_norm_policy,

498 out_q=out_q_prefill,

499 + out_k=out_k,

500 + out_v=out_v,

501 )

502 attn_metadata.hpc_prefill_q_scale = q_scale

503 else:

504 hpc.rope_norm_store_kv(

359 - key_cache,

360 - value_cache,

505 + kcache_arg,

506 + vcache_arg,

507 qkv_prefill,

508 self.cos_sin_cache,

509 seq_lens_prefill,

510 cu_seqlens_prefill,

365 - block_table_prefill,

511 + kv_indices_arg,

512 True, # is_prefill

513 q_norm_weight=q_norm_weight,

514 k_norm_weight=k_norm_weight,

515 out_q=out_q_prefill,

516 + out_k=out_k,

517 + out_v=out_v,

518 qk_norm_policy=self.qk_norm_policy,

519 )

520

521 + if out_k is not None:

522 + self._zero_pad_last_blocks(

523 + key_cache, value_cache, seq_lens_prefill, block_table_prefill

524 + )

525 + prefill_slots = slot_mapping[num_decode_tokens:num_actual_tokens]

526 + self._scatter_kv_cache(key_cache, out_k, prefill_slots)

527 + self._scatter_kv_cache(value_cache, out_v, prefill_slots)

528 +

529 # --- Decode ---

530 if num_decode_reqs > 0:

531 num_seq_kvcache = attn_metadata.seq_lens[:num_decode_reqs]

@@ -380,15 +536,32 @@ class HpcRopeNorm(CustomOp, HpcModule):

536 decode_query_len = attn_metadata.decode_query_len

537 out_q_decode = output[:num_decode_tokens]

538

539 + if fused_kv_write:

540 + out_k = out_v = None

541 + kcache_arg, vcache_arg = key_cache, value_cache

542 + kv_indices_arg = block_table_decode

543 + else:

544 + out_k = torch.empty(

545 + num_decode_tokens,

546 + self.num_kv_heads,

547 + self.head_dim,

548 + dtype=key_cache.dtype,

549 + device=qkv.device,

550 + )

551 + out_v = torch.empty_like(out_k)

552 + scratch = self._kv_write_scratch(key_cache)

553 + kcache_arg, vcache_arg = scratch[0], scratch[1]

554 + kv_indices_arg = torch.zeros_like(block_table_decode)

555 +

556 if self.use_fp8:

557 _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(

385 - key_cache=key_cache,

386 - value_cache=value_cache,

558 + key_cache=kcache_arg,

559 + value_cache=vcache_arg,

560 qkv=qkv_decode,

561 cos_sin=self.cos_sin_cache,

562 num_seqlen_per_req=num_seq_kvcache,

563 q_index=attn_metadata.qo_indptr_decode,

391 - kvcache_indices=block_table_decode,

564 + kvcache_indices=kv_indices_arg,

565 is_prefill=False,

566 k_scale=k_scale,

567 v_scale=v_scale,

@@ -398,26 +571,38 @@ class HpcRopeNorm(CustomOp, HpcModule):

571 k_norm_weight=k_norm_weight,

572 qk_norm_policy=self.qk_norm_policy,

573 out_q=out_q_decode,

574 + out_k=out_k,

575 + out_v=out_v,

576 )

577 attn_metadata.hpc_decode_q_scale = q_scale

578 if split_k_flag is not None:

579 attn_metadata.hpc_split_k_flag = split_k_flag

580 else:

581 hpc.rope_norm_store_kv(

407 - key_cache,

408 - value_cache,

582 + kcache_arg,

583 + vcache_arg,

584 qkv_decode,

585 self.cos_sin_cache,

586 num_seq_kvcache,

587 attn_metadata.qo_indptr_decode,

413 - block_table_decode,

588 + kv_indices_arg,

589 False, # is_prefill

590 q_norm_weight=q_norm_weight,

591 k_norm_weight=k_norm_weight,

592 out_q=out_q_decode,

593 + out_k=out_k,

594 + out_v=out_v,

595 qk_norm_policy=self.qk_norm_policy,

596 )

597

598 + if out_k is not None:

599 + self._zero_pad_last_blocks(

600 + key_cache, value_cache, num_seq_kvcache, block_table_decode

601 + )

602 + decode_slots = slot_mapping[:num_decode_tokens]

603 + self._scatter_kv_cache(key_cache, out_k, decode_slots)

604 + self._scatter_kv_cache(value_cache, out_v, decode_slots)

605 +

606 # Signal HpcAttentionImpl that KV cache has already been written by

607 # rope_norm_store_kv[_fp8] above, so it should skip its own

608 # reshape_and_cache_flash. Set after the fused kernels ran (either

[AttentionBackend][HPC-ops] update hpc rope norm to support stride kv cache (#53705)