feat: SGLang-style paged attention kernels replace page-table path

- PagedAttentionParams uses flat KV pool + req_to_token + kv_indptr/qo_indptr instead of page_table
- MMA split-KV decode and split-Q prefill kernels with indirect ragged-batch addressing
- Prefill kernel accepts 4D mask (causal-aware); decode kernel supports 2D mask
- CudaBackend is inference-only: kv_cache=None raises, no torch fallback
- benchmark.py: required --ckpt, --backend/--compare options
- Parallel build isolates build-temp/build-lib per subprocess
- Standalone test covers decode/prefill with mask, 27 cases pass
This commit is contained in:
2026-08-01 15:41:25 +08:00
parent 9960f79920
commit 41dcf0feb9
17 changed files with 1683 additions and 535 deletions
+14 -6
View File
@@ -13,18 +13,26 @@ from tests.extension.conftest import D, skip_no_kernel
@skip_no_kernel
def test_training_forward_matches_torch(cuda_model):
"""Training forward (kv_cache=None) should produce identical logits."""
"""Training forward (kv_cache=None) should produce identical logits.
CudaBackend is inference-only: it raises when kv_cache is None. Training
must use TorchNativeBackend (the default). Verify the torch path is
stable and that CudaBackend rejects the training path explicitly.
"""
import pytest
model, _ = cuda_model
input_ids = torch.randint(0, 1000, (2, 16), device="cuda")
with torch.no_grad():
out_torch = model(input_ids)
with attn_backend(ATTN_BACKEND.CUDA):
with torch.no_grad():
out_cuda = model(input_ids)
diff = (out_torch["logits"].float() - out_cuda["logits"].float()).abs().max().item()
assert diff == 0.0, f"Training forward diff {diff} should be 0"
with pytest.raises(RuntimeError, match="does not support training"):
with attn_backend(ATTN_BACKEND.CUDA):
with torch.no_grad():
model(input_ids)
assert out_torch["logits"].shape[0] == 2
@skip_no_kernel