feat: primus-turbo attn add sbhd format support

[RuibinCheung]

• Primus Turbo attention add sbhd format support. Eliminate extra transpose kernel call in Attention.

[Copilot]

Pull request overview

Adds experimental support for additional QKV tensor layouts (notably sbhd) in the Primus Turbo attention wrapper, while introducing special-casing for sink attention to force a specific layout.

Changes:

• Removes the previous manual sbhd -> bshd transpose and instead forwards qkv_format into the underlying flash_attn op.
• Introduces a use_sink_attn flag and forces sink-attention execution to use bshd, including explicit tensor permutations for Q/K/V and the output.

[Copilot]

When use_sink_attn is enabled, query/key/value are always permuted as if the incoming layout were sbhd (S,B,H,D) -> bshd (B,S,H,D). If qkv_format is already bshd (or any non-sbhd value coming from packed_seq_params), this permutation will corrupt the tensor layout while qkv_format is forced to "bshd", creating a format/tensor mismatch.

Consider either (a) explicitly asserting qkv_format == "sbhd" before permuting in the sink-attention path, or (b) permuting conditionally based on qkv_format and ensuring o is permuted back consistently to preserve the expected output layout.

[Copilot]

qkv_format is now passed through to self.attn(...) without any local validation. Since this value can come from packed_seq_params, an unexpected value will likely fail deeper in the kernel with a less actionable error (and could also interact badly with the explicit format conversions in the sink-attention path).

Recommend validating qkv_format against the set of supported formats in this backend and raising a clear ValueError (or keeping an assert) before using it.

Suggested change

	qkv_format = packed_seq_kwargs.get("qkv_format", self.qkv_format)
	qkv_format = packed_seq_kwargs.get("qkv_format", self.qkv_format)
	supported_qkv_formats = ("sbhd", "bshd", "thd")
	if qkv_format not in supported_qkv_formats:
	raise ValueError(
	f"Unsupported qkv_format: {qkv_format}. "
	f"Supported formats: {supported_qkv_formats}"
	)

[Copilot]

Typo/grammar in the new comment: "sink attention only support bshd format" → "sink attention only supports bshd format".

Suggested change

	# NOTE: sink attention only support bshd format
	# NOTE: sink attention only supports bshd format

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

[Copilot]

qkv_format is now only handled for the "sbhd" case; any other value falls through and is treated as if inputs are already BSHD, which can silently produce wrong results if qkv_format is misspelled/unsupported. Please add explicit validation (e.g., accept only sbhd and bshd, and raise ValueError otherwise) rather than relying on implicit fallthrough or assert.

Suggested change

	value = value.permute(1, 0, 2, 3)
	value = value.permute(1, 0, 2, 3)
	elif qkv_format == "bshd":
	pass
	else:
	raise ValueError(
	f"Unsupported qkv_format: {qkv_format!r}. Expected one of ('sbhd', 'bshd')."
	)

[Copilot]

This adds conditional layout handling based on qkv_format, but there are no tests covering the new sbhd/bshd behavior. Please add a focused unit test that sets qkv_format to both values and asserts the returned tensor layout/shape matches expectations.