Skip to content
Open
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions lightx2v_platform/ops/attn/hygon_dcu/flash_attn.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ def half(x):
softmax_scale = 1.0 / math.sqrt(q.shape[-1])
# Use Flash Attention 2.6.1 (ROCm version) with varlen interface
if SAPRDE_LINEAR_ATTN and int(os.getenv("USE_SLA", 0)) and q.shape[1] == k.shape[1]:
topk_value = float(os.getenv("SPARSE_ATTN_TOPK", "0.5"))
topk_value = float(os.getenv("SPARSE_ATTN_TOPK", "0.4"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Parsing the environment variable SPARSE_ATTN_TOPK directly into a float without validation can lead to runtime crashes (if the value is not a valid float) or unexpected behavior/crashes in the underlying sparse attention kernel (if the value is outside the valid range of (0.0, 1.0]). It is safer to wrap this in a try-except block and validate that the resulting value is within the expected range.

            try:
                topk_value = float(os.getenv("SPARSE_ATTN_TOPK", "0.4"))
                if not (0.0 < topk_value <= 1.0):
                    raise ValueError
            except ValueError:
                logger.warning("Invalid SPARSE_ATTN_TOPK value. Falling back to default 0.4.")
                topk_value = 0.4


q = q_flat.unsqueeze(0)
k = k_flat.unsqueeze(0)
Expand All @@ -124,7 +124,7 @@ def half(x):
q,
k,
v,
topk=0.4,
topk=topk_value,
)
else:
output = flash_attn_varlen_func(
Expand Down