Repository navigation
Conversation
The VSA metadata builder and route builder capped their shape caches at 16 entries and failed the request that introduced the 17th shape. The CUDA Graph runner that captures the cached tensor addresses has no such cap, so the limit only turned a new resolution/frame profile into a runtime error. Drop the cap: both caches keep one entry per distinct shape for as long as the graphs that reference them and are released together with the graphs in cleanup. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
The Triton block pooling kernel loaded tokens in fixed-width steps and only masked against the sequence length, so a block size that is not a multiple of the step width read into the next block. Mask the block boundary as well, reject non-positive block sizes, and check the kernel against the PyTorch fallback for block sizes below, not a multiple of, and equal to the step width. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
… checks Add the VSA and SOL predictor kernel suites to the CPU and B200 test lists. Move the TRTLLM skip-softmax CUDA Graph capture test out of the cpu_only module into the TRTLLM metadata suite with an SM100 gate so it runs on the GPU stage. Assert that TRTLLM VSA executes the block-sparse FMHA rather than only predicting routes, check the forwarded VSA gate values, verify that the Wan VSA pipeline reference takes the SDPA fallback, and reject non-finite tensor timesteps. Annotate the helpers and tests flagged during review and read the sparse algorithm discriminator directly instead of through getattr. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Owner
Author
|
Superseded: the same changes landed as one commit (85b6b4d) directly on |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Follow-up for NVIDIA#18079 that resolves the open review threads there. It targets the PR branch
user/yuhangh/update_vsa, so merging it updates NVIDIA#18079 in place.Threads resolved by this PR
l0_cpu.ymlwith +1 from BowenFu (thread)test_attention_vsa_kernels.pyandtest_sol_predictor_kernels.pyare listed inl0_b200.yml; the SOL kernel suite is also listed inl0_cpu.ymlfor itscpu_onlytests.pipeline_wan.py(thread)cleanup. No new user-facing knob.block_poolblock sizes (thread)PrimsTSBlockSparseFmha.forwardrunning instead of a non-Noneprediction.test_skip_softmax.py(thread)test_trtllm_attention_metadata.py(B200 list) under an SM100/SM103 gate. Inside thecpu_onlymodule it never ran in CI.inf, tensornanand tensorinfcases added.test_wan_vsa_pipeline.py(thread)to_sparse_paramsreturn type (thread) and the four nitpicks (_compact_qkv,_predict,getattr, quoted annotation)Threads already addressed by the current NVIDIA#18079 tip
SOLSparsePredictor(coderabbit): the predictor is a stateless functional operator now.pipeline_ltx2_two_stages.pyquestion (lfr-0531): the PR no longer touches the LTX-2 model or pipeline code; main carries the timestep threading via [TRTLLM-15917][feat] Integrate Sol-Attn sparse attention into VisualGen NVIDIA/TensorRT-LLM#18329, and the stubbed tests were removed.dense_layers(coderabbit): the field islist[int]now.fmha_cachesprecondition key (coderabbit): already fixed in the refined commits.AttentionConfigvalidates the backend/algorithm pairing; the model only checks the algorithm.Test
B200 container, branch based on the current NVIDIA#18079 tip.
pre-commit run --files <changed files>clean.test_timestep_phase.py,test_sol_predictor_kernels.py,test_skip_softmax.py,test_trtllm_attention_metadata.py,test_attention_vsa.py,test_sol_attention.py,test_visual_gen_args.py, the three multi-GPU unit tests): 218 passed, 41 skipped.test_sol_predictor_kernels.py,test_trtllm_attention_metadata.pyincluding the moved capture test,test_attention_vsa_kernels.py,test_attention_vsa.py): 136 passed.test_attention_integration.py -k vsa: 4 passed, including the block-sparse FMHA execution check.test_wan_vsa_pipeline.py::TestWanVsa14B_PipelineCorrectness::test_cosine_similaritywithLLM_MODELS_ROOTset: 1 passed; the CuTe run launched the CuTe fine stage and the SDPA reference did not.