From facc6509e748e732c36c1a65cbbaecbee36576d5 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 19 May 2026 17:26:50 +0000 Subject: [PATCH] Fix imports in vLLM codepaths test --- tests/test_vllm_codepaths_b200.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_vllm_codepaths_b200.py b/tests/test_vllm_codepaths_b200.py index 711553f1..346d3aa1 100644 --- a/tests/test_vllm_codepaths_b200.py +++ b/tests/test_vllm_codepaths_b200.py @@ -98,11 +98,10 @@ def causal_prefill_attention(q, kv, scale): def main(): """Test the exact csa_attention.py code paths used in the container.""" from cutedsl.blackwell_attention import ( - blackwell_attention_kv_write, - blackwell_attention_decode, - blackwell_attention_forward, + apply_gptj_rope, + apply_inv_gptj_rope, ) - # Also import the vLLM patch version + # Import the vLLM patch version (the actual code used in the container) sys.path.insert(0, os.path.join(REPO, "vllm", "patches", "layers")) from csa_attention import ( fused_qnorm_rope_kv_insert_py, @@ -110,6 +109,7 @@ def main(): blackwell_attention_decode as vllm_decode, kv_quantize_fp8 as vllm_kv_quantize, kv_dequantize_fp8 as vllm_kv_dequantize, + causal_prefill_attention, ) torch.cuda.set_device(0)