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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user