[perf] Parallelize apply_token_bitmask across vocab tiles

[vuuihc]

Motivation

apply_token_bitmask is used by SGLang structured output to set disallowed token logits to -inf.
The current AscendC implementation partitions work only across rows:

blockDim = min(num_rows, core_num);

Structured decoding commonly has a small effective batch and a large vocabulary. With batch=1,
the current kernel therefore launches one AIV even for 128K-200K vocabularies.

Change

This PR flattens the independent (row, vocab tile) work space and distributes it with a grid-stride loop:

total_work = num_rows * num_tiles
work_idx   = block_idx + k * block_dim
row_id     = work_idx / num_tiles
tile_id    = work_idx % num_tiles

The host chooses an aligned tile size that exposes enough work to cover the available AIVs while respecting
the UB-derived maximum tile size. The implementation uses runtime AIV and UB properties and does not contain
an Ascend 950-specific capability gate.

The production diff is limited to:

• host-side tile sizing and launch geometry;
• kernel-side flattened work distribution;
• correctness coverage for a realistic single-row 151936-vocabulary shape and reliable category test exits.

The scalar bit-unpacking loop is unchanged and remains a separate optimization opportunity.

Correctness and compatibility

Work items write disjoint (row, vocab tile) regions, so no cross-core synchronization is required.
The production implementation was validated separately from the experiment-only legacy op:

• Ascend 950PR production build: passed
• Boundary correctness: 6/6
• LLM correctness: 24/24
• General correctness: 21/21
• Ascend910_9382 (A3) compile validation: passed
• Ascend910B1 (A2) compile validation: passed
• Local pre-commit: passed

The Ascend 950 runtime validation used the latest #632 build baseline because the production PR intentionally
does not include unrelated Ascend 950 build changes. The production implementation itself has no SOC-specific
branch.

Production validation summary

Performance

Environment:

• Device: Ascend950PR_9579
• Compile target: Ascend950PR_9599
• CANN: 9.0.0
• PyTorch: 2.7.1+cpu
• torch_npu: 2.7.1.post4
• Timing: torch.npu.Event(enable_timing=True)
• 20 warmups and 100 measured iterations, alternating legacy/candidate launch order
• Same wheel and input tensors for legacy/candidate comparison
• Clean tracked source state recorded in JSON (git_status == "")

Case	Legacy P50 (ms)	Optimized P50 (ms)	P50 speedup	P99 speedup
b1, v32K, FP16, sparse	0.237408	0.012190	19.476x	18.771x
b1, v128256, FP16, sparse	0.947396	0.025066	37.796x	37.185x
b1, v151936, BF16, sparse	1.156653	0.035527	32.557x	31.157x
b8, v151936, BF16, sparse	1.166488	0.182100	6.406x	6.376x
b1, v200000, BF16, sparse	1.520744	0.041408	36.726x	34.934x
b1, v151936, BF16, random	1.356192	0.039714	34.149x	33.595x
b1, v151936, BF16, all unmasked	0.083114	0.016077	5.170x	3.111x
b2, v151936, BF16, sparse	1.160851	0.061702	18.814x	17.832x
b32, v151936, BF16, sparse	1.179318	0.868482	1.358x	1.359x
b1, v32100, FP16, sparse	0.244726	0.017299	14.147x	13.577x
indices=8, b16, v151936, BF16	1.220371	0.235937	5.172x	5.113x

Geometric-mean P50 speedup across all 11 cases: 13.044x. No case regressed.
The speedup decreases as batch grows, matching the expected mechanism: row-only scheduling already exposes
more parallelism at larger batch sizes.

Artifacts:

• Summary
• Raw samples and metadata
• Functional test log
• Hardware and software snapshot

Related work

• feat: add apply_token_bitmask operator #442 introduced the AscendC apply_token_bitmask operator.
• sglang#24133 connected SGLang structured output to this operator.
• [feat] support Ascend950PR_9599 SOC_VERSION kernel compile #632 provides the Ascend 950 kernel build baseline used only by the validation branch.

cc @CarterDuan @ChefWu551 @iforgetmyname

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.