biondizzle
f3b551956d
- Move dead dsv4/ modules to dsv4/_archive/ (52 files)
- model/{dsv4,mtp,layer,layer_schedule}
- layers/{embedding,attention,ffn,norm} (kept linear,mhc,router,moe,shared_expert,grouped_linear - live)
- cache/*, kernels/cache/*, kernels/indexer/{csa_indexer,score_topk,compute_valid_lens}
- kernels/router/{nvfp4_fused_router,dense_router_decode_kernel,dense_router_prefill}
- ops/{topk,topk_select,rope,router}, loader/{hf_checkpoint,layout_convert}
- reference/{attention,compressor,csa_attention,moe_pipeline}
- kernels/compressor/{compress_tail,csa_hca}
- Restore dsv4/ops/{router,custom_ops}.py (needed by live layers)
- Fix dsv4/kernels/{indexer,compressor,attention}/__init__.py (removed broken imports)
- Remove preload_all() from loader.py (dead, referenced nonexistent .cu file)
- Fix loader.py docstring (fused_amax_quantize_nvfp4 → quantize_nvfp4_from_buffer)
- Move broken tests to tests/e2e_archive/
- test_fused_router, production_values_test, e2e/{one_layer,model_construction,csa_hca}
- vLLM has 0 imports of dsv4 (Step 0 confirmed)
52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
"""Python wrapper for the append_swa CUDA kernel.
|
|
|
|
Writes raw BF16 KV into the FP8/BF16 split state cache layout.
|
|
Quantizes the non-RoPE half BF16 -> FP8 (E4M3 amax-based scaling),
|
|
writes the RoPE half as-is, computes per-token inverse scale, and
|
|
updates the ring buffer head + position field.
|
|
|
|
One block per token. Threads cooperatively:
|
|
1. Compute amax over fp8-dim elements (warp reduce).
|
|
2. Quantize BF16 -> FP8 with per-token scale.
|
|
3. Write FP8 entries + BF16 RoPE entries + inv_scale + position.
|
|
4. Atomic increment ring buffer head.
|
|
"""
|
|
import os
|
|
import torch
|
|
from torch.utils.cpp_extension import load
|
|
|
|
_kernel_module = None
|
|
|
|
|
|
def _get_kernel_module():
|
|
global _kernel_module
|
|
if _kernel_module is not None:
|
|
return _kernel_module
|
|
kernel_dir = os.path.join(os.path.dirname(__file__), "..", "cuda")
|
|
_kernel_module = load(
|
|
name="append_swa",
|
|
sources=[os.path.join(kernel_dir, "append_swa.cu")],
|
|
extra_cuda_cflags=["-O3", "--generate-code=arch=compute_100a,code=[sm_100a]"],
|
|
verbose=False,
|
|
)
|
|
return _kernel_module
|
|
|
|
|
|
def append_swa_kernel(
|
|
raw_kv: torch.Tensor, # (T, head_dim) BF16
|
|
request_slots: torch.Tensor, # (T,) int32
|
|
positions: torch.Tensor, # (T,) int32
|
|
swa_fp8: torch.Tensor, # (max_req, n_win, fp8_dim) uint8
|
|
swa_rope: torch.Tensor, # (max_req, n_win, rope_dim) BF16
|
|
swa_inv: torch.Tensor, # (max_req, n_win) FP32
|
|
swa_pos: torch.Tensor, # (max_req, n_win) int32
|
|
swa_head: torch.Tensor, # (max_req,) int32
|
|
rope_dim: int,
|
|
):
|
|
mod = _get_kernel_module()
|
|
mod.append_swa(
|
|
raw_kv, request_slots, positions,
|
|
swa_fp8, swa_rope, swa_inv, swa_pos, swa_head,
|
|
rope_dim,
|
|
)
|